perf(router): skip unused deployment reads in the callback filter step and cache anthropic capability probes (#50641)

litellm/router_rust/steps/callbacks.py

@@ -452,6 +452,10 @@ class OptionalChecks:
             *self._policies_not_set_globally(),
         )
 
+    def has_filters(self) -> bool:
+        """Whether ``filter_deployments`` would run any callback: without one it keeps every deployment."""
+        return any(filters_deployments(callback) for callback in self.routing_callbacks())
+
     async def filter_deployments(
         self,
         model: str,
@@ -464,12 +468,10 @@ class OptionalChecks:
         the Rust core takes only if its later filters keep it; a filter's error is raised for the
         Rust core to decide on."""
         remaining: list[object] = list(deployments)  # mutable-ok: filter contract  # rebind-ok: narrowed
-        base: Final[object] = CustomLogger.async_filter_deployments  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # untyped deployment dicts
         request_kwargs: Final = dict(request)
         callbacks: Final = tuple(self.routing_callbacks())
         for callback in callbacks:
-            # The base filter returns its input unchanged: only overriding filters are run.
-            if not isinstance(callback, CustomLogger) or type(callback).async_filter_deployments is base:  # pyright: ignore[reportUnknownMemberType]  # see above
+            if not isinstance(callback, CustomLogger) or not filters_deployments(callback):
                 continue
             if isinstance(callback, PromptCachingDeploymentCheck):
                 continue
@@ -572,14 +574,22 @@ def _ids_of(value: object) -> Sequence[object]:
     return value if isinstance(value, list | tuple) else ()  # pyright: ignore[reportUnknownVariableType]  # the step's id list
 
 
+def filters_deployments(callback: object) -> bool:
+    """A routing callback that overrides ``async_filter_deployments``: the base filter returns its
+    input unchanged, so only overriding filters are run."""
+    base: Final[object] = CustomLogger.async_filter_deployments  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # untyped deployment dicts
+    return isinstance(callback, CustomLogger) and type(callback).async_filter_deployments is not base  # pyright: ignore[reportUnknownMemberType]  # see above
+
+
 def filtered_ids(remaining: Iterable[object], candidates: Sequence[str]) -> list[str]:  # mutable-ok: a JSON list
     """The ids of the deployments the filters left, in their order, among the candidates."""
-    ids: Final[list[str]] = []  # mutable-ok: built here
+    allowed: Final = frozenset(candidates)
+    ids: Final[dict[str, None]] = {}  # mutable-ok: built here, insertion-ordered
     for deployment in remaining:
         found = (as_str_mapping((as_str_mapping(deployment) or {}).get("model_info")) or {}).get("id")
-        if isinstance(found, str) and found in candidates and found not in ids:
-            ids.append(found)
-    return ids
+        if isinstance(found, str) and found in allowed:
+            ids.setdefault(found)
+    return list(ids)
 
 
 class RustOptionalChecks:
@@ -698,6 +708,8 @@ class RustOptionalChecks:
         from litellm.router_rust.core import attempt_failure
 
         candidates: Final = [str(id) for id in step.get("candidates") or ()]  # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType, reportUnknownArgumentType]  # the step's id list
+        if not self._opt.has_filters():
+            return (list(dict.fromkeys(candidates)), None), None
         records: Final = self._records_of(candidates)  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType]  # RustCore's deployment dicts
         narrowed: Final = _request_with_failed_ids(request, step)
         try:

litellm/rust_bridge/host/model_capabilities.py

@@ -1,6 +1,27 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
+from functools import lru_cache
+from typing import Final, TypeAlias
+
 import litellm
+from litellm.litellm_core_utils.core_helpers import as_str_mapping
+
+
+class _Held:
+    __slots__ = ("value",)
+
+    def __init__(self, value: object) -> None:
+        self.value = value
+
+    def __hash__(self) -> int:
+        return id(self.value)
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, _Held) and other.value is self.value
+
+
+_Fingerprint: TypeAlias = tuple[str, str | None, _Held, int, int, tuple[_Held, ...]]
 
 
 def _resolved_provider(model: str, custom_llm_provider: str | None) -> tuple[str, str]:
@@ -11,11 +32,50 @@ def _resolved_provider(model: str, custom_llm_provider: str | None) -> tuple[str
     return resolved_model, provider
 
 
+def _fingerprint(model: str, custom_llm_provider: str | None, resolved_model: str, provider: str) -> _Fingerprint:
+    """Changes whenever an answer could: ``register_model`` bumps the generation, and replacing
+    any entry the probes look up (``litellm.model_cost[key] = ...``) changes that entry's identity.
+    The key holds the entries it names, so a cached answer's entries are never freed and their
+    addresses never reused."""
+    from litellm.llms.anthropic.common_utils import AnthropicModelInfo
+    from litellm.utils import get_model_cost_mutation_generation
+
+    model_cost: Final = as_str_mapping(litellm.model_cost) or {}  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]  # the model map is an untyped dict of rows
+    keys: Final = dict.fromkeys(
+        (
+            *AnthropicModelInfo._model_map_lookup_candidates(model),  # pyright: ignore[reportPrivateUsage]  # the keys the capability walk reads
+            resolved_model,
+            f"{provider}/{resolved_model}",
+        )
+    )
+    return (
+        model,
+        custom_llm_provider,
+        _Held(model_cost),
+        get_model_cost_mutation_generation(),
+        len(model_cost),
+        tuple(_Held(model_cost.get(key)) for key in keys),
+    )
+
+
 def anthropic_model_capabilities(model: str, custom_llm_provider: str | None) -> dict[str, object]:
+    """The capability probes the Rust messages route reads, once per model-map state: each probe
+    walks the model map (``get_model_info`` plus generalization matching), and the route asks on
+    every request."""
+    resolved_model, provider = _resolved_provider(model, custom_llm_provider)
+    cached: Final = _capabilities(
+        _fingerprint(model, custom_llm_provider, resolved_model, provider), resolved_model, provider
+    )
+    tiers: Final = as_str_mapping(cached["effort_tiers"])
+    return {**cached, "effort_tiers": dict(tiers) if tiers is not None else cached["effort_tiers"]}
+
+
+@lru_cache(maxsize=256)
+def _capabilities(fingerprint: _Fingerprint, resolved_model: str, provider: str) -> Mapping[str, object]:
     from litellm.llms.anthropic.chat.transformation import AnthropicConfig
     from litellm.llms.anthropic.common_utils import AnthropicModelInfo
 
-    resolved_model, provider = _resolved_provider(model, custom_llm_provider)
+    model: Final = fingerprint[0]
 
     def supports(flag: str) -> bool:
         return AnthropicModelInfo._supports_model_capability(model, flag, provider)  # pyright: ignore[reportPrivateUsage]  # same probes the Python transform runs; forking them would drift

tests/unit/rust_bridge/host/test_model_capabilities.py

@@ -86,3 +86,72 @@ def test_capability_source_observes_runtime_registration_changes(monkeypatch: py
 def test_default_thinking_capability_is_projected(model: str, expected: bool) -> None:
     capabilities: Final = anthropic_model_capabilities(model, "anthropic")
     assert capabilities["thinking_on_by_default"] is expected
+
+
+def test_capabilities_are_probed_once_per_model_map_state(monkeypatch: pytest.MonkeyPatch) -> None:
+    from litellm.llms.anthropic.common_utils import AnthropicModelInfo
+
+    model: Final = "claude-test-probe-count"
+    _flag_model(monkeypatch, model, supports_reasoning=True)
+    probes: list[str] = []
+    probe: Final = AnthropicModelInfo._supports_model_capability
+
+    def counted(name: str, key: str, provider: str) -> bool:
+        probes.append(key)
+        return probe(name, key, provider)
+
+    monkeypatch.setattr(AnthropicModelInfo, "_supports_model_capability", staticmethod(counted))
+
+    first: Final = anthropic_model_capabilities(model, "anthropic")
+    probed: Final = len(probes)
+    second: Final = anthropic_model_capabilities(model, "anthropic")
+
+    assert probed > 0
+    assert len(probes) == probed
+    assert second == first
+
+
+def test_a_cached_answer_is_a_fresh_copy(monkeypatch: pytest.MonkeyPatch) -> None:
+    model: Final = "claude-test-fresh-copy"
+    _flag_model(monkeypatch, model, supports_xhigh_reasoning_effort=True)
+    first: Final = anthropic_model_capabilities(model, "anthropic")
+    tiers: Final = first["effort_tiers"]
+    assert isinstance(tiers, dict)
+    first["supports_reasoning"] = "changed"
+    tiers["xhigh"] = "changed"
+
+    second: Final = anthropic_model_capabilities(model, "anthropic")
+
+    assert second["supports_reasoning"] is False
+    assert second["effort_tiers"] == {
+        **dict.fromkeys(("minimal", "low", "medium", "high", "max"), False),
+        "xhigh": True,
+    }
+
+
+def test_register_model_is_observed_after_a_cached_answer() -> None:
+    model: Final = "claude-test-registered-later"
+    try:
+        litellm.register_model({model: {"litellm_provider": "anthropic", "mode": "chat", "supports_reasoning": False}})
+        before: Final = anthropic_model_capabilities(model, "anthropic")
+        litellm.register_model({model: {"litellm_provider": "anthropic", "mode": "chat", "supports_reasoning": True}})
+        after: Final = anthropic_model_capabilities(model, "anthropic")
+    finally:
+        litellm.model_cost.pop(model, None)
+
+    assert before["supports_reasoning"] is False
+    assert after["supports_reasoning"] is True
+
+
+def test_the_cache_key_holds_the_entries_it_names(monkeypatch: pytest.MonkeyPatch) -> None:
+    from litellm.rust_bridge.host.model_capabilities import _fingerprint  # pyright: ignore[reportPrivateUsage]
+
+    model: Final = "claude-test-held-entry"
+    _flag_model(monkeypatch, model, supports_reasoning=True)
+    entry: Final = litellm.model_cost[model]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # the model map is an untyped dict of rows
+    held: Final = _fingerprint(model, "anthropic", model, "anthropic")
+    _flag_model(monkeypatch, model, supports_reasoning=True)
+    same_content: Final = _fingerprint(model, "anthropic", model, "anthropic")
+
+    assert any(h.value is entry for h in held[5])
+    assert held != same_content

tests/unit/test_router/core/test_callback_filters_step.py

@@ -0,0 +1,84 @@
+"""
+The ``callback_filters`` host step (``RustCallbackChecks._answer_callback_filters``,
+litellm/router_rust/steps/callbacks.py): a pick reads the candidates' deployments back only when a
+routing callback filters them, and the kept ids come back in order, once each.
+"""
+
+from collections.abc import Mapping
+from typing import Final
+
+import pytest
+
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.router_rust.steps.callbacks import filtered_ids, filters_deployments
+from tests.unit.test_router.core.registry_helpers import MakeCore
+from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment
+
+GROUP: Final = [deployment("g", "a"), deployment("g", "b"), deployment("g", "c")]
+
+
+class Narrow(CustomLogger):
+    """Keeps the deployments other than ``dropped``."""
+
+    def __init__(self, dropped: str) -> None:
+        super().__init__()
+        self.dropped = dropped
+
+    async def async_filter_deployments(
+        self,
+        model: str,
+        healthy_deployments: list,
+        messages: object,
+        request_kwargs: object = None,
+        parent_otel_span: object = None,
+    ) -> list:
+        return [d for d in healthy_deployments if d["model_info"]["id"] != self.dropped]
+
+
+class Observer(CustomLogger):
+    """A callback that does not override ``async_filter_deployments``."""
+
+
+def _entry(deployment_id: object) -> Mapping[str, object]:
+    return {"model_info": {"id": deployment_id}}
+
+
+def test_filtered_ids_keep_the_filter_order_once_each_among_the_candidates() -> None:
+    remaining: Final = [_entry("c"), _entry("x"), _entry("a"), _entry("c"), _entry(7), {"model_info": None}, "junk"]
+
+    assert filtered_ids(remaining, ["a", "b", "c"]) == ["c", "a"]
+
+
+def test_only_an_overriding_filter_counts_as_a_filter() -> None:
+    assert filters_deployments(Narrow(dropped="a"))
+    assert not filters_deployments(Observer())
+    assert not filters_deployments(object())
+
+
+@pytest.mark.asyncio
+async def test_a_pick_without_filters_reads_no_deployment_back(
+    make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    host.lists["callbacks"].append(Observer())
+    core: Final = make_core(model_list=list(GROUP))
+
+    def refuse(_self: object, _ids: object) -> list[dict[str, object]]:  # mutable-ok: the patched signature
+        raise AssertionError("no filter runs, so no deployment is read back")
+
+    monkeypatch.setattr(type(core), "_records_of", refuse)
+
+    picked: Final = {(await core.select("g"))["model_info"]["id"] for _ in range(30)}  # pyright: ignore[reportIndexIssue]
+
+    assert picked == {"a", "b", "c"}
+
+
+@pytest.mark.asyncio
+async def test_a_pick_with_a_filter_reads_the_candidates_and_honours_it(
+    make_core: MakeCore, host: ScriptedHost
+) -> None:
+    host.lists["callbacks"].append(Narrow(dropped="b"))
+    core: Final = make_core(model_list=list(GROUP))
+
+    picked: Final = {(await core.select("g"))["model_info"]["id"] for _ in range(30)}  # pyright: ignore[reportIndexIssue]
+
+    assert picked == {"a", "c"}