From cd7306913052e53f646cd9437daaaf475c5f7021 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 00:39:29 +0000 Subject: [PATCH] feat(router): choose the router backend through a Route.ROUTER catalog rule Router(...) now resolves Route.ROUTER once at construction. PYTHON builds PythonRouter as before. Under RUST_OPT_IN (LITELLM_RUST) the Rust backend is built when it can serve the config and the call declines to PythonRouter otherwise; RUST_REQUIRED raises naming the unsupported argument, deployment field or global setting. rust_support decides what the first increment can serve: the reliability arguments (retries, fallbacks, retry policies, cooldowns, Redis) on simple-shuffle, without wildcards, strategy routers, orders, tags, tpm/rpm limits, team deployments, aliases or callbacks that filter deployments. RustRouter is a stub until the native router exists, so every config still declines today. Also types the facade's conditional base as a TypeAlias and marks its forwarding kwargs. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Aw1xktpUnix5tWhvSzey8t --- litellm/router.py | 17 +- litellm/router_backends/rust_router.py | 36 +++++ litellm/router_backends/rust_support.py | 151 ++++++++++++++++++ litellm/router_backends/selection.py | 50 ++++++ litellm/rust_bridge/catalog.py | 2 + tests/unit/router_backends/__init__.py | 0 .../unit/router_backends/test_rust_support.py | 91 +++++++++++ tests/unit/router_backends/test_selection.py | 46 ++++++ tests/unit/rust_bridge/test_catalog.py | 2 +- 9 files changed, 387 insertions(+), 8 deletions(-) create mode 100644 litellm/router_backends/rust_router.py create mode 100644 litellm/router_backends/rust_support.py create mode 100644 litellm/router_backends/selection.py create mode 100644 tests/unit/router_backends/__init__.py create mode 100644 tests/unit/router_backends/test_rust_support.py create mode 100644 tests/unit/router_backends/test_selection.py diff --git a/litellm/router.py b/litellm/router.py index fbf58ac5594..be029dbd089 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1,7 +1,9 @@ import inspect -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Final, TypeAlias, cast from litellm.router_backends.python_router import PythonRouter +from litellm.router_backends.rust_router import RustRouter +from litellm.router_backends.selection import select_backend from litellm.types.router import ( AlertingConfig, AllowedFailsPolicy, @@ -28,9 +30,9 @@ __all__ = ( # Provisional (router-poc.md, "Facade surface"): Router exposes the backend's full surface, # raw internals included, until callers move to a typed protocol. if TYPE_CHECKING: - _RouterSurface = PythonRouter + _RouterSurface: TypeAlias = PythonRouter else: - _RouterSurface = object + _RouterSurface: TypeAlias = object class _Forwarded: @@ -40,7 +42,8 @@ class _Forwarded: self._name = name def __get__(self, instance: "Router | None", owner: type) -> object: - return cast(object, getattr(PythonRouter if instance is None else instance.backend, self._name)) # cast-ok: forwarded attribute + target: Final = PythonRouter if instance is None else instance.backend + return cast(object, getattr(target, self._name)) # cast-ok: forwarded attribute def __set__(self, instance: "Router", value: object) -> None: setattr(instance.backend, self._name, value) @@ -50,12 +53,12 @@ class _Forwarded: class Router(_RouterSurface): - backend: PythonRouter + backend: PythonRouter | RustRouter if not TYPE_CHECKING: - def __init__(self, *args, **kwargs): - object.__setattr__(self, "backend", PythonRouter(*args, **kwargs)) + def __init__(self, *args, **kwargs): # kwargs-ok: forwards PythonRouter.__init__'s signature + object.__setattr__(self, "backend", select_backend(args, kwargs)) __init__.__signature__ = inspect.signature(PythonRouter.__init__) diff --git a/litellm/router_backends/rust_router.py b/litellm/router_backends/rust_router.py new file mode 100644 index 00000000000..0e28766a14b --- /dev/null +++ b/litellm/router_backends/rust_router.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +from litellm.router_backends.rust_support import unsupported_reason +from litellm.rust_bridge.bindings import NativeBinding + + +@dataclass(frozen=True, slots=True) +class RustRouterDeclined: + reason: str + + +def _native_router_type(value: object) -> type | None: + return value if isinstance(value, type) else None + + +NATIVE_ROUTER: Final = NativeBinding("Router", validate=_native_router_type) + + +class RustRouter: + """The Rust-backed router. Members it does not emulate yet raise `NotImplementedError` naming them.""" + + def __getattr__(self, name: str) -> object: + raise NotImplementedError(f"the Rust router backend does not emulate Router.{name} yet") + + +def build_rust_router(arguments: Mapping[str, object]) -> RustRouter | RustRouterDeclined: + reason: Final = unsupported_reason(arguments) + if reason is not None: + return RustRouterDeclined(reason) + if NATIVE_ROUTER.load() is None: + return RustRouterDeclined("the native router is not available in this build") + return RustRouterDeclined("the native router is not wired up yet") diff --git a/litellm/router_backends/rust_support.py b/litellm/router_backends/rust_support.py new file mode 100644 index 00000000000..e103442f0aa --- /dev/null +++ b/litellm/router_backends/rust_support.py @@ -0,0 +1,151 @@ +"""What the Rust router backend can serve today, decided once per `Router(...)` construction.""" + +from __future__ import annotations + +import inspect +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final, cast + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_backends.python_router import PythonRouter +from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model + +_INIT: Final = cast(Callable[..., None], PythonRouter.__init__) # cast-ok: only the signature is read +_SIGNATURE: Final = inspect.signature(_INIT) + + +def _default(parameter: inspect.Parameter) -> object: + return cast(object, parameter.default) # cast-ok: inspect types defaults as Any + + +_DEFAULTS: Final[Mapping[str, object]] = MappingProxyType( + {name: _default(parameter) for name, parameter in _SIGNATURE.parameters.items()} +) +_CALLBACKS: Final = TypeAdapter(tuple[object, ...]) + +SUPPORTED_ARGUMENTS: Final = frozenset( + { + "model_list", + "redis_url", + "redis_host", + "redis_port", + "redis_password", + "redis_db", + "num_retries", + "max_fallbacks", + "timeout", + "stream_timeout", + "set_verbose", + "debug_level", + "default_fallbacks", + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "retry_after", + "retry_policy", + "model_group_retry_policy", + "allowed_fails", + "cooldown_time", + "disable_cooldowns", + } +) + +_UNSUPPORTED_DEPLOYMENT_KEYS: Final = frozenset({"tpm", "rpm", "max_parallel_requests"}) +_UNSUPPORTED_LITELLM_PARAMS: Final = frozenset( + {"order", "tags", "tag_regex", "tpm", "rpm", "max_parallel_requests", "silent_model"} +) +_UNSUPPORTED_MODEL_INFO: Final = frozenset({"team_id", "order", "allowed_fails", "allowed_fails_policy", "blocked"}) +_CALLBACK_HOOKS: Final = ("async_filter_deployments", "async_pre_call_check") + + +class _RawDeployment(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + model_name: str + litellm_params: Mapping[str, object] + model_info: Mapping[str, object] | None = None + + +_MODEL_LIST: Final = TypeAdapter(tuple[_RawDeployment, ...]) + + +def bind_arguments(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Mapping[str, object]: + """The explicitly passed `Router(...)` arguments by parameter name.""" + bound: Final = _SIGNATURE.bind(None, *args, **kwargs) + arguments: Final = cast(Mapping[str, object], bound.arguments) # cast-ok: inspect types bound values as Any + return {name: value for name, value in arguments.items() if name != "self"} + + +def unsupported_reason(arguments: Mapping[str, object]) -> str | None: + """Why the Rust backend cannot serve this configuration, or None when it can.""" + return next( + ( + reason + for reason in ( + _argument_reason(arguments), + _model_list_reason(arguments.get("model_list")), + _global_reason(), + ) + if reason is not None + ), + None, + ) + + +def _argument_reason(arguments: Mapping[str, object]) -> str | None: + return next( + ( + f"Router argument {name!r} is not supported by the Rust router yet" + for name, value in arguments.items() + if name not in SUPPORTED_ARGUMENTS and value != _DEFAULTS[name] + ), + None, + ) + + +def _model_list_reason(model_list: object) -> str | None: + if model_list is None: + return None + try: + deployments: Final = _MODEL_LIST.validate_python(model_list) + except ValidationError: + return "model_list entries the Rust router cannot read" + return next( + (reason for reason in map(_deployment_reason, deployments) if reason is not None), + None, + ) + + +def _deployment_reason(deployment: _RawDeployment) -> str | None: + extra: Final = deployment.model_extra or {} + model: Final = deployment.litellm_params.get("model") + keys: Final = ( + *(f"{key}" for key in _UNSUPPORTED_DEPLOYMENT_KEYS.intersection(extra)), + *(f"litellm_params.{key}" for key in _UNSUPPORTED_LITELLM_PARAMS.intersection(deployment.litellm_params)), + *(f"model_info.{key}" for key in _UNSUPPORTED_MODEL_INFO.intersection(deployment.model_info or {})), + ) + if "*" in deployment.model_name: + return f"wildcard model_name {deployment.model_name!r}" + if isinstance(model, str) and classify_strategy_router_model(model) is not None: + return f"strategy router deployment {model!r}" + if keys: + return f"deployment {deployment.model_name!r} sets {', '.join(sorted(keys))}" + return None + + +def _global_reason() -> str | None: + if litellm.model_alias_map: + return "litellm.model_alias_map is set" + hooked: Final = tuple( + type(callback).__name__ + for callback in _CALLBACKS.validate_python(cast(object, litellm.callbacks)) # cast-ok: untyped global + if isinstance(callback, CustomLogger) + and any(getattr(type(callback), hook) is not getattr(CustomLogger, hook) for hook in _CALLBACK_HOOKS) + ) + if hooked: + return f"callbacks {', '.join(hooked)} filter or pre-check deployments" + return None diff --git a/litellm/router_backends/selection.py b/litellm/router_backends/selection.py new file mode 100644 index 00000000000..28aa4b055d2 --- /dev/null +++ b/litellm/router_backends/selection.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Final, TypeAlias, cast + +from typing_extensions import assert_never + +from litellm._logging import verbose_router_logger +from litellm.router_backends.python_router import PythonRouter +from litellm.router_backends.rust_router import RustRouter, RustRouterDeclined, build_rust_router +from litellm.router_backends.rust_support import bind_arguments +from litellm.rust_bridge.catalog import Route, RouteContext, Rules, decision +from litellm.rust_bridge.configuration import Decision + +_PYTHON_ROUTER_TYPE: TypeAlias = Callable[..., PythonRouter] +_PYTHON_ROUTER: Final = cast(_PYTHON_ROUTER_TYPE, PythonRouter) # cast-ok: Router(...) forwards untyped arguments + + +class RustRouterUnsupportedError(ValueError): + pass + + +def select_backend( + args: tuple[object, ...], + kwargs: Mapping[str, object], + rules: Rules | None = None, +) -> PythonRouter | RustRouter: + selected: Final = decision(RouteContext(Route.ROUTER), rules) + match selected: + case Decision.PYTHON: + return _PYTHON_ROUTER(*args, **kwargs) + case Decision.RUST_WITH_FALLBACK | Decision.RUST_REQUIRED: + built: Final = build_rust_router(bind_arguments(args, kwargs)) + if isinstance(built, RustRouter): + return built + return _declined(built, selected, args, kwargs) + case _: + assert_never(selected) + + +def _declined( + declined: RustRouterDeclined, + selected: Decision, + args: tuple[object, ...], + kwargs: Mapping[str, object], +) -> PythonRouter: + if selected is Decision.RUST_REQUIRED: + raise RustRouterUnsupportedError(f"the Rust router is required but cannot serve this config: {declined.reason}") + verbose_router_logger.info("Rust router declined, using the Python router: %s", declined.reason) + return _PYTHON_ROUTER(*args, **kwargs) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 73e95c24de4..7ac8a2adcd6 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -23,6 +23,7 @@ class Route(str, Enum): OCR = "ocr" TOKEN_COUNTER = "token_counter" TOKENIZER = "tokenizer" + ROUTER = "router" @dataclass(frozen=True, slots=True) @@ -76,6 +77,7 @@ RULES: Final[Rules] = ( RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), RouteRule(Route.TOKENIZER, Rollout.PYTHON_ONLY), RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), + RouteRule(Route.ROUTER, Rollout.RUST_OPT_IN), ) diff --git a/tests/unit/router_backends/__init__.py b/tests/unit/router_backends/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_backends/test_rust_support.py b/tests/unit/router_backends/test_rust_support.py new file mode 100644 index 00000000000..c6b1f1a7285 --- /dev/null +++ b/tests/unit/router_backends/test_rust_support.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_backends.rust_support import bind_arguments, unsupported_reason + + +def _deployment(overrides: Mapping[str, object] = MappingProxyType({})) -> Mapping[str, object]: + return {"model_name": "gpt", "litellm_params": {"model": "openai/gpt-4o-mini"}, **overrides} + + +@pytest.mark.parametrize( + "arguments", + ( + {}, + {"model_list": [_deployment()]}, + {"model_list": [_deployment(), _deployment({"litellm_params": {"model": "openai/gpt-4o", "weight": 2}})]}, + { + "model_list": [_deployment({"model_info": {"id": "a", "cooldown_time": 3}})], + "num_retries": 3, + "fallbacks": [{"gpt": ["other"]}], + "context_window_fallbacks": [{"gpt": ["big"]}], + "content_policy_fallbacks": [{"gpt": ["safe"]}], + "retry_policy": {"RateLimitErrorRetries": 2}, + "model_group_retry_policy": {"gpt": {"TimeoutErrorRetries": 1}}, + "allowed_fails": 2, + "cooldown_time": 9, + "disable_cooldowns": False, + "redis_host": "localhost", + "redis_port": 6379, + }, + {"routing_strategy": "simple-shuffle", "enable_pre_call_checks": False, "fallbacks": []}, + ), +) +def test_supported_configs_have_no_reason(arguments: Mapping[str, object]) -> None: + assert unsupported_reason(arguments) is None + + +@pytest.mark.parametrize( + ("arguments", "named"), + ( + ({"routing_strategy": "latency-based-routing"}, "'routing_strategy'"), + ({"enable_pre_call_checks": True}, "'enable_pre_call_checks'"), + ({"model_group_alias": {"a": "gpt"}}, "'model_group_alias'"), + ({"model_list": [_deployment({"model_name": "openai/*"})]}, "wildcard"), + ({"model_list": [_deployment({"litellm_params": {"model": "auto_router/x"}})]}, "strategy router"), + ({"model_list": [_deployment({"litellm_params": {"model": "m", "order": 1}})]}, "litellm_params.order"), + ({"model_list": [_deployment({"litellm_params": {"model": "m", "rpm": 10}})]}, "litellm_params.rpm"), + ({"model_list": [_deployment({"model_info": {"team_id": "t"}})]}, "model_info.team_id"), + ({"model_list": [_deployment({"tpm": 100})]}, "tpm"), + ({"model_list": [{"litellm_params": {"model": "m"}}]}, "cannot read"), + ), +) +def test_unsupported_configs_name_the_reason(arguments: Mapping[str, object], named: str) -> None: + reason: Final = unsupported_reason(arguments) + + assert reason is not None + assert named in reason + + +def test_global_model_alias_map_is_unsupported(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {"a": "gpt"}) + + assert unsupported_reason({}) == "litellm.model_alias_map is set" + + +async def _filter_nothing(*_: object) -> tuple[()]: + return () + + +_Filtering: Final = type("_Filtering", (CustomLogger,), {"async_filter_deployments": _filter_nothing}) + + +def test_callbacks_that_filter_deployments_are_unsupported_but_plain_loggers_are_not( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "callbacks", [CustomLogger()]) + assert unsupported_reason({}) is None + + monkeypatch.setattr(litellm, "callbacks", [CustomLogger(), _Filtering()]) + assert unsupported_reason({}) == "callbacks _Filtering filter or pre-check deployments" + + +def test_bind_arguments_maps_positional_and_keyword_arguments_by_name() -> None: + assert bind_arguments(([_deployment()],), {"num_retries": 1}) == {"model_list": [_deployment()], "num_retries": 1} diff --git a/tests/unit/router_backends/test_selection.py b/tests/unit/router_backends/test_selection.py new file mode 100644 index 00000000000..eea7d3e831f --- /dev/null +++ b/tests/unit/router_backends/test_selection.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from collections.abc import Generator +from typing import Final + +import pytest + +from litellm.router_backends.python_router import PythonRouter +from litellm.router_backends.selection import RustRouterUnsupportedError, select_backend +from litellm.rust_bridge import configuration +from litellm.rust_bridge.catalog import Route, RouteRule +from litellm.rust_bridge.configuration import Rollout + +MODEL_LIST: Final = ({"model_name": "gpt", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}},) + + +@pytest.fixture(autouse=True) +def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: + monkeypatch.delenv("LITELLM_RUST", raising=False) + configuration.reset_rust_configuration() + yield + configuration.reset_rust_configuration() + + +def test_python_decision_builds_the_python_router_with_the_callers_arguments() -> None: + backend: Final = select_backend((), {"model_list": list(MODEL_LIST), "num_retries": 7}, rules=()) + + assert isinstance(backend, PythonRouter) + assert backend.num_retries == 7 + assert backend.get_model_names() == ["gpt"] + + +def test_opt_in_declines_an_unsupported_config_to_the_python_router() -> None: + rules: Final = (RouteRule(Route.ROUTER, Rollout.RUST_OPT_OUT),) + + backend: Final = select_backend((), {"model_list": list(MODEL_LIST), "routing_strategy": "least-busy"}, rules) + + assert isinstance(backend, PythonRouter) + assert backend.routing_strategy == "least-busy" + + +def test_required_raises_naming_what_the_rust_router_cannot_serve() -> None: + rules: Final = (RouteRule(Route.ROUTER, Rollout.RUST_REQUIRED),) + + with pytest.raises(RustRouterUnsupportedError, match="'routing_strategy'"): + select_backend((), {"model_list": list(MODEL_LIST), "routing_strategy": "least-busy"}, rules) diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index 4242294a33a..46ddb2383c2 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -45,7 +45,7 @@ def test_shipped_decisions( if route is Route.OCR or (route is Route.TRANSCRIPTION and provider == "bedrock"): assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED - elif route is Route.MESSAGES and provider == "anthropic": + elif (route is Route.MESSAGES and provider == "anthropic") or route is Route.ROUTER: assert catalog.rollout(context) is Rollout.RUST_OPT_IN opted_in: Final = environment == "1" or (environment is None and process is True) assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if opted_in else Decision.PYTHON)