From edb511ef37ae16e9108f08003f1ac502af75cc2b Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 21:42:57 +0000 Subject: [PATCH] feat(router): serve the proxy's reads from rust and hand over to python for the rest Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Aw1xktpUnix5tWhvSzey8t --- litellm/router_backends/rust_call.py | 21 +- litellm/router_backends/rust_router.py | 265 ++++++++++++++---- litellm/router_backends/rust_support.py | 13 +- litellm/router_backends/rust_surface.py | 88 ++++++ .../router/test_proxy_switchover.py | 124 ++++++++ .../router/test_redis_parity.py | 28 ++ .../router/test_rust_router.py | 4 +- 7 files changed, 468 insertions(+), 75 deletions(-) create mode 100644 litellm/router_backends/rust_surface.py diff --git a/litellm/router_backends/rust_call.py b/litellm/router_backends/rust_call.py index 5b61f3ac365..e38c187879a 100644 --- a/litellm/router_backends/rust_call.py +++ b/litellm/router_backends/rust_call.py @@ -14,14 +14,16 @@ import copy import time from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence from types import MappingProxyType -from typing import Final, Literal, Protocol, TypeAlias, cast +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, cast from pydantic import BaseModel, ConfigDict, TypeAdapter import litellm from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_guardrail import is_guardrail_intervention -from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.asyncify import ( + run_async_function, # pyright: ignore[reportUnknownVariableType] # untyped helper, wrapped in _run_sync +) from litellm.litellm_core_utils.exception_mapping_utils import ( _get_response_headers, # pyright: ignore[reportPrivateUsage] # the header reader the cooldown callback uses ) @@ -72,6 +74,9 @@ from litellm.utils import ( _get_retry_after_from_exception_header, # pyright: ignore[reportPrivateUsage] # the Retry-After parser retries and cooldowns use ) +if TYPE_CHECKING: + from litellm.router import Router + Bucket: TypeAlias = MutableMapping[str, object] Kwargs: TypeAlias = dict[str, object] Operation: TypeAlias = Literal["completion", "responses", "anthropic_messages"] @@ -382,15 +387,14 @@ class RoutedCall: "model": parsed.model, self._metadata_key: self._buckets[parsed.bucket], } - router: Final = self._normalizer + router: Final = cast("Router", self._normalizer) # cast-ok: the PythonRouter the helpers read checks off return await _is_fallback_target_authorized( router, parsed.target, parsed.model_group, kwargs ) and await _is_fallback_target_within_budget(router, parsed.target, parsed.model_group, kwargs) def allow_fallback_sync(self, check: Mapping[str, object]) -> bool: """Python's sync calls run the async fallback loop, checks included, through `run_async_function`.""" - allowed: Final[object] = run_async_function(self.allow_fallback, check) - return allowed is True + return _run_sync(lambda: self.allow_fallback(check)) is True def success(self, response: object, outcome: Mapping[str, object], ops: Sequence[Mapping[str, object]]) -> object: self._apply(ops) @@ -841,6 +845,13 @@ class RoutedCall: error.message = f"{message}{text}" # pyright: ignore[reportAttributeAccessIssue] # litellm exceptions carry a message +def _run_sync(start: Callable[[], Awaitable[object]]) -> object: + run: Final = cast( + Callable[[Callable[[], Awaitable[object]]], object], run_async_function + ) # cast-ok: untyped helper + return run(start) + + def _handler(operation: Operation) -> Callable[..., Awaitable[object]]: """Read at call time, so a patched handler is the one called.""" name: Final = _GENERIC_HANDLERS[operation] diff --git a/litellm/router_backends/rust_router.py b/litellm/router_backends/rust_router.py index 87022007fc4..13b76c5b10c 100644 --- a/litellm/router_backends/rust_router.py +++ b/litellm/router_backends/rust_router.py @@ -8,11 +8,12 @@ from types import MappingProxyType from typing import Final, Protocol, TypeAlias, TypeVar, cast from urllib.parse import quote -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, TypeAdapter import litellm from litellm import constants from litellm._logging import verbose_router_logger +from litellm.caching.redis_cache import RedisCache from litellm.constants import RUNTIME_UPDATABLE_ROUTER_SETTINGS from litellm.router_backends.python_router import PythonRouter from litellm.router_backends.rust_call import ( @@ -25,6 +26,7 @@ from litellm.router_backends.rust_call import ( metadata_key_for, ) from litellm.router_backends.rust_support import unsupported_reason +from litellm.router_backends.rust_surface import ASSIGNABLE, VIEWS from litellm.router_utils import clientside_credential_handler from litellm.router_utils.cooldown_handlers import ( _is_allowed_fails_set_on_router, # pyright: ignore[reportPrivateUsage] # the rule PythonRouter's cooldowns apply @@ -68,6 +70,8 @@ _UNSUPPORTED_REQUEST_KWARGS: Final = ( "model_group_retry_policy", "include_fallback_errors", "_router_weights", + "routing_strategy", + "enable_tag_filtering", ) _MOCK_FAILURES: Final = ( ("mock_testing_fallbacks", "fallbacks"), @@ -124,11 +128,37 @@ class NormalizerView(AttemptRouter, Protocol): def get_settings(self) -> Mapping[str, object]: ... + def _update_redis_cache(self, cache: RedisCache) -> None: ... + def _update_kwargs_before_fallbacks( self, model: str, kwargs: Kwargs, metadata_variable_name: str = "metadata" ) -> None: ... +class PythonServing(Protocol): + """The `PythonRouter` methods a request the Rust router does not serve is forwarded to, typed.""" + + async def acompletion( + self, model: str, messages: Sequence[Mapping[str, object]], **kwargs: object + ) -> object: ... # kwargs-ok: forwards Router.acompletion's kwargs + + def completion( + self, model: str, messages: Sequence[Mapping[str, object]], **kwargs: object + ) -> object: ... # kwargs-ok: forwards Router.completion's kwargs + + async def aresponses( + self, custom_llm_provider: str | None = None, client: object = None, **kwargs: object + ) -> object: ... # kwargs-ok: forwards Router.aresponses' kwargs + + async def aanthropic_messages( + self, custom_llm_provider: str | None = None, client: object = None, **kwargs: object + ) -> object: ... # kwargs-ok: forwards Router.aanthropic_messages' kwargs + + +def _serving(python: PythonRouter) -> PythonServing: + return cast(PythonServing, python) # cast-ok: PythonRouter's methods take untyped kwargs + + def _raw(deployment: Deployment) -> Mapping[str, object]: return deployment.model_dump(exclude_none=True) @@ -213,24 +243,61 @@ def _providers() -> tuple[str, ...]: def redis_url(arguments: Mapping[str, object]) -> str | None: """The URL `PythonRouter` would connect its router cache to, if any.""" - url: Final = arguments.get("redis_url") + return _url( + arguments.get("redis_url"), + arguments.get("redis_host"), + arguments.get("redis_port"), + arguments.get("redis_password"), + arguments.get("redis_db"), + ) + + +_PLAIN_REDIS_KWARGS: Final = frozenset({"url", "host", "port", "username", "password", "db"}) +_CLIENT_TUNING_REDIS_KWARGS: Final = frozenset( + {"socket_timeout", "socket_connect_timeout", "health_check_interval", "retry_on_timeout", "max_connections"} +) +_REDIS_KWARGS: Final = TypeAdapter(Mapping[str, object]) + + +def adopted_redis_url(cache: RedisCache) -> str | None: + """The URL of a `RedisCache` the proxy hands the router, or None when it connects some other + way (cluster, sentinel, TLS) that the native store does not open. Client tuning (timeouts, + pool size) is left to the native store's own defaults.""" + configured: Final = _REDIS_KWARGS.validate_python( + cast(object, cache.redis_kwargs) + ) # cast-ok: RedisCache leaves it untyped + kwargs: Final = {key: value for key, value in configured.items() if value is not None} + if not kwargs.keys() <= _PLAIN_REDIS_KWARGS | _CLIENT_TUNING_REDIS_KWARGS: + return None + return _url( + kwargs.get("url"), + kwargs.get("host"), + kwargs.get("port"), + kwargs.get("password"), + kwargs.get("db"), + kwargs.get("username"), + ) + + +def _url(url: object, host: object, port: object, password: object, db: object, username: object = None) -> str | None: if isinstance(url, str): return url - host: Final = arguments.get("redis_host") - port: Final = arguments.get("redis_port") if host is None or port is None: return None - password: Final = arguments.get("redis_password") - credentials: Final = f":{quote(str(password), safe='')}@" if password else "" - return f"redis://{credentials}{host}:{port}/{arguments.get('redis_db') or 0}" + user: Final = quote(str(username), safe="") if username else "" + credentials: Final = f"{user}:{quote(str(password), safe='')}@" if password else (f"{user}@" if user else "") + return f"redis://{credentials}{host}:{port}/{db or 0}" class RustRouter: - """The Rust-backed router. Members it does not emulate yet raise `NotImplementedError` naming them. + """The Rust-backed router. Deployments are normalized by a `PythonRouter` built from the same arguments with its router - callbacks removed (`discard`), which also builds each attempt's kwargs. Routing decisions, - cooldowns and usage counters belong to the native router. + callbacks removed (`discard`), which also builds each attempt's kwargs and answers the reads + in `rust_surface.VIEWS`. Routing decisions, cooldowns and usage counters belong to the native + router. Anything else it is asked for, a member, a request option or a runtime change it does + not serve, hands the instance over to a `PythonRouter` for good (`on_fall_back`), or raises + `RustRouterUnsupportedError` when Rust is required. """ def __init__( @@ -249,59 +316,126 @@ class RustRouter: object.__setattr__(self, "_updates", MappingProxyType({})) object.__setattr__(self, "_required", required) object.__setattr__(self, "_fall_back", None) - object.__setattr__( - self, - "_native", - native( - [_project(deployment) for deployment in normalizer.model_list], - _settings(normalizer), - list(_providers()), - redis_url(arguments), - seed, - ), - ) + object.__setattr__(self, "_handed_to", None) + object.__setattr__(self, "_assigned", MappingProxyType({})) + object.__setattr__(self, "_redis_cache", None) + object.__setattr__(self, "_native_type", native) + object.__setattr__(self, "_seed", seed) + object.__setattr__(self, "_native", self._build_native(redis_url(arguments))) _normalizer: NormalizerView _native: NativeRouter + _native_type: NativeRouterType + _seed: int | None _arguments: Mapping[str, object] _updates: Mapping[str, object] + _assigned: Mapping[str, object] + _redis_cache: RedisCache | None _required: bool _fall_back: FallBack | None + _handed_to: PythonRouter | None def __getattr__(self, name: str) -> object: normalizer: Final[object] = self.__dict__.get("_normalizer") + handed_to: Final[object] = self.__dict__.get("_handed_to") + if handed_to is not None: + return _attribute(handed_to, name) if normalizer is None or not hasattr(normalizer, name): raise AttributeError(name) - raise NotImplementedError(f"the Rust router backend does not emulate Router.{name} yet") + if name in VIEWS: + return _attribute(normalizer, name) + return _attribute(self._python(f"Router.{name} is not served by the Rust router"), name) def __setattr__(self, name: str, value: object) -> None: - raise NotImplementedError(f"the Rust router backend does not emulate setting Router.{name} yet") + if self._handed_to is None and name in ASSIGNABLE: + setattr(self._normalizer, name, value) + object.__setattr__(self, "_assigned", MappingProxyType({**self._assigned, name: value})) + return + setattr(self._python(f"setting Router.{name} is not served by the Rust router"), name, value) + + def _build_native(self, url: str | None) -> NativeRouter: + return self._native_type( + [_project(deployment) for deployment in self._normalizer.model_list], + _settings(self._normalizer), + list(_providers()), + url, + self._seed, + ) + + def _python(self, reason: str) -> PythonRouter: + """The `PythonRouter` this instance has handed over to, built from its current model list, + settings, assigned attributes and Redis on the first request for one. In-memory state is + not carried over; Redis state is.""" + if self._handed_to is not None: + return self._handed_to + if self._required or self._fall_back is None: + raise RustRouterUnsupportedError(f"the Rust router cannot serve this: {reason}") + python: Final = _PYTHON_ROUTER(**{**self._arguments, "model_list": copy.deepcopy(self._normalizer.model_list)}) + view: Final = cast(NormalizerView, python) # cast-ok: PythonRouter's untyped update_settings + if self._updates: + view.update_settings(**self._updates) + if self._redis_cache is not None: + view._update_redis_cache(cache=self._redis_cache) # pyright: ignore[reportPrivateUsage] # the proxy's own hand-off + for name, value in self._assigned.items(): + setattr(python, name, value) + verbose_router_logger.info("Rust router handing over to the Python router for good: %s", reason) + object.__setattr__(self, "_handed_to", python) + self._fall_back(python) + return python @property def model_list(self) -> list[dict[str, object]]: # mutable-ok: PythonRouter's read surface - return self._normalizer.model_list + return self._reader.model_list @property def model_names(self) -> set[str]: # mutable-ok: PythonRouter's read surface - return self._normalizer.model_names + return self._reader.model_names @property def total_calls(self) -> Mapping[str, int]: - return self._normalizer.total_calls + return self._reader.total_calls @property def success_calls(self) -> Mapping[str, int]: - return self._normalizer.success_calls + return self._reader.success_calls @property def fail_calls(self) -> Mapping[str, int]: - return self._normalizer.fail_calls + return self._reader.fail_calls def get_model_names(self) -> list[str]: # mutable-ok: PythonRouter's read surface - return self._normalizer.get_model_names() + return self._reader.get_model_names() + + @property + def _reader(self) -> NormalizerView: + return ( + cast(NormalizerView, self._handed_to) if self._handed_to is not None else self._normalizer + ) # cast-ok: a PythonRouter def discard(self) -> None: - return None + if self._handed_to is not None: + self._handed_to.discard() + + def arm_routing_read_prefetch(self, model: str, request_kwargs: Mapping[str, object] | None = None) -> None: + """Python's routing reads batch their Redis lookups through this; the native router reads its own.""" + if self._handed_to is not None: + self._handed_to.arm_routing_read_prefetch(model, dict(request_kwargs or {})) + + def _update_redis_cache(self, cache: RedisCache) -> None: + """The proxy's Redis, attached after construction: the native router reconnects to it, which + drops its in-memory state (the proxy attaches it at startup), and the Python side attaches it + too, so its cooldown reads see what the native router writes.""" + if self._handed_to is not None: + self._handed_to._update_redis_cache(cache) # pyright: ignore[reportPrivateUsage] # the proxy's own hand-off + return + url: Final = adopted_redis_url(cache) + if url is None: + python: Final = self._python("the Rust router cannot open this Redis connection") + python._update_redis_cache(cache) # pyright: ignore[reportPrivateUsage] # the proxy's own hand-off + return + self._normalizer._update_redis_cache(cache) # pyright: ignore[reportPrivateUsage] # the proxy's own hand-off + object.__setattr__(self, "_redis_cache", cache) + object.__setattr__(self, "_native", self._build_native(url)) def on_fall_back(self, fall_back: FallBack) -> None: """Where this instance hands over to a `PythonRouter` when a runtime change asks for something @@ -309,34 +443,36 @@ class RustRouter: object.__setattr__(self, "_fall_back", fall_back) def upsert_deployment(self, deployment: Deployment) -> Deployment | None: - return self._change(lambda: self._normalizer.upsert_deployment(deployment), deployments=(deployment,)) + return self._change(lambda router: router.upsert_deployment(deployment), deployments=(deployment,)) def add_deployment(self, deployment: Deployment) -> Deployment | None: - return self._change(lambda: self._normalizer.add_deployment(deployment), deployments=(deployment,)) + return self._change(lambda router: router.add_deployment(deployment), deployments=(deployment,)) def delete_deployment(self, id: str) -> Deployment | None: - return self._change(lambda: self._normalizer.delete_deployment(id)) + return self._change(lambda router: router.delete_deployment(id)) def update_settings(self, **kwargs: object) -> None: # kwargs-ok: PythonRouter.update_settings' surface updatable: Final = {name: value for name, value in kwargs.items() if name in RUNTIME_UPDATABLE_ROUTER_SETTINGS} - self._change(lambda: self._normalizer.update_settings(**kwargs), updates=updatable) + self._change(lambda router: router.update_settings(**kwargs), updates=updatable) @property def retry_policy(self) -> object: - return self._normalizer.retry_policy + return self._reader.retry_policy def get_settings(self) -> Mapping[str, object]: - return self._normalizer.get_settings() + return self._reader.get_settings() def _change( self, - apply: Callable[[], _T], + apply: Callable[[NormalizerView], _T], deployments: Sequence[Deployment] = (), updates: Mapping[str, object] = MappingProxyType({}), ) -> _T: """Applies a runtime change on the Python side, which normalizes deployments and resolves - settings, then swaps the native snapshot. A change the Rust backend cannot serve falls back - to `PythonRouter` for good, or raises when Rust is required.""" + settings, then swaps the native snapshot. A change the Rust backend cannot serve is applied + to the `PythonRouter` this instance hands over to.""" + if self._handed_to is not None: + return apply(self._reader) merged: Final = MappingProxyType({**self._updates, **updates}) prospective: Final = { **self._arguments, @@ -345,32 +481,24 @@ class RustRouter: } reason: Final = unsupported_reason(prospective) if reason is not None: - return self._hand_over(reason, apply, merged) - result: Final = apply() + self._python(f"a runtime change asks for {reason}") + return apply(self._reader) + result: Final = apply(self._normalizer) object.__setattr__(self, "_updates", merged) self._native.replace( [_project(deployment) for deployment in self._normalizer.model_list], _settings(self._normalizer) ) return result - def _hand_over(self, reason: str, apply: Callable[[], _T], updates: Mapping[str, object]) -> _T: - if self._required or self._fall_back is None: - raise RustRouterUnsupportedError(f"the Rust router cannot serve this change: {reason}") - result: Final = apply() - python: Final = _PYTHON_ROUTER(**{**self._arguments, "model_list": copy.deepcopy(self._normalizer.model_list)}) - if updates: - settings_view: Final = cast(NormalizerView, python) # cast-ok: PythonRouter's untyped update_settings - settings_view.update_settings(**updates) - verbose_router_logger.info("Rust router handing over to the Python router for good: %s", reason) - self._fall_back(python) - return result - async def acompletion( self, model: str, messages: Sequence[Mapping[str, object]], **kwargs: object, # kwargs-ok: forwards litellm.acompletion's kwargs ) -> object: + python: Final = self._python_for(kwargs) + if python is not None: + return await _serving(python).acompletion(model, messages, **kwargs) call: Final = self._call("completion", model, {**kwargs, "messages": messages}) return await self._route(call, True) @@ -380,6 +508,9 @@ class RustRouter: messages: Sequence[Mapping[str, object]], **kwargs: object, # kwargs-ok: forwards litellm.completion's kwargs ) -> object: + python: Final = self._python_for(kwargs) + if python is not None: + return _serving(python).completion(model, messages, **kwargs) call: Final = self._call("completion", model, {**kwargs, "messages": messages}) return self._native.route(call.spec, call.driver, False) @@ -390,7 +521,9 @@ class RustRouter: **kwargs: object, # kwargs-ok: forwards litellm.aresponses' kwargs ) -> object: """`factory_function`'s wrapper drops its own `custom_llm_provider` and `client`.""" - del custom_llm_provider, client + python: Final = self._python_for(kwargs) + if python is not None: + return await _serving(python).aresponses(custom_llm_provider=custom_llm_provider, client=client, **kwargs) return await self._route(self._call("responses", _model(kwargs), kwargs), True) async def aanthropic_messages( @@ -400,21 +533,31 @@ class RustRouter: **kwargs: object, # kwargs-ok: forwards litellm.anthropic_messages' kwargs ) -> object: """`factory_function`'s wrapper drops its own `custom_llm_provider` and `client`.""" - del custom_llm_provider, client + python: Final = self._python_for(kwargs) + if python is not None: + return await _serving(python).aanthropic_messages( + custom_llm_provider=custom_llm_provider, client=client, **kwargs + ) return await self._route(self._call("anthropic_messages", _model(kwargs), kwargs), True) + def _python_for(self, kwargs: Kwargs) -> PythonRouter | None: + """The `PythonRouter` to serve this request instead: the one this instance handed over to, + or a new one when the request asks for something the Rust router does not serve.""" + if self._handed_to is not None: + return self._handed_to + unsupported: Final = next((name for name in _UNSUPPORTED_REQUEST_KWARGS if name in kwargs), None) + if unsupported is not None: + return self._python(f"request option {unsupported!r} is not served by the Rust router") + if _IS_CLIENTSIDE_CREDENTIAL(kwargs): + return self._python("client-side credentials are not served by the Rust router") + return None + async def _route(self, call: _Call, asynchronous: bool) -> object: native_route: Final = self._native.route(call.spec, call.driver, asynchronous) routed: Final = cast(Awaitable[object], native_route) # cast-ok: an asynchronous route is a coroutine return await routed def _call(self, operation: Operation, model: str, kwargs: Kwargs) -> _Call: - unsupported: Final = [name for name in _UNSUPPORTED_REQUEST_KWARGS if name in kwargs] - if unsupported: - raise NotImplementedError(f"the Rust router backend does not support request option {unsupported[0]!r} yet") - clientside: Final = _IS_CLIENTSIDE_CREDENTIAL(kwargs) - if clientside: - raise NotImplementedError("the Rust router backend does not support client-side credentials yet") kwargs["model"] = model # rebind-ok: PythonRouter stamps the caller's kwargs the same way if operation == "completion": kwargs.setdefault("stream", False) diff --git a/litellm/router_backends/rust_support.py b/litellm/router_backends/rust_support.py index 3819c5cb29f..5104c57a23b 100644 --- a/litellm/router_backends/rust_support.py +++ b/litellm/router_backends/rust_support.py @@ -140,15 +140,15 @@ def _fallbacks_reason(arguments: Mapping[str, object]) -> str | None: def _group_names(model_list: object) -> frozenset[str]: return frozenset( - deployment.model_name for deployment in map(_readable, _entries(model_list)) if deployment is not None + deployment.model_name for deployment in map(_readable, _entries(model_list) or ()) if deployment is not None ) -def _entries(model_list: object) -> tuple[object, ...]: +def _entries(model_list: object) -> tuple[object, ...] | None: try: return _MODEL_LIST.validate_python(model_list or ()) except ValidationError: - return () + return None def _readable(entry: object) -> _RawDeployment | None: @@ -181,11 +181,10 @@ def _plain_chain_of(key: str, targets: object, groups: frozenset[str]) -> bool: def _model_list_reason(model_list: object, ignore_invalid: bool) -> str | None: """With `ignore_invalid_deployments`, Python drops an entry it cannot read instead of failing, so the Rust router only serves the ones it can.""" - if model_list is None: - return None - if not isinstance(model_list, (list, tuple)): + entries: Final = _entries(model_list) + if entries is None: return "a model_list the Rust router cannot read" - deployments: Final = tuple(map(_readable, _entries(model_list))) + deployments: Final = tuple(map(_readable, entries)) if None in deployments and not ignore_invalid: return "model_list entries the Rust router cannot read" return next( diff --git a/litellm/router_backends/rust_surface.py b/litellm/router_backends/rust_surface.py new file mode 100644 index 00000000000..710a0c4346e --- /dev/null +++ b/litellm/router_backends/rust_surface.py @@ -0,0 +1,88 @@ +"""Which `Router` members the Rust backend answers without handing the instance over to Python. + +Requests it routes itself are methods on `RustRouter`. The names below are answered by the +Python-side router it keeps for normalization, which runtime changes keep in step with the +native snapshot, so its model list, settings and config-derived lookups are the ones Rust +routes by. Cooldown reads there see Rust's cooldowns once the router shares the proxy's Redis. +Any other member hands the instance over to `PythonRouter` for good. +""" + +from typing import Final + +VIEWS: Final = frozenset( + { + "_are_all_deployments_blocked", + "_get_client", + "_get_model_from_alias", + "_get_router_deployment_budget_limiter", + "_is_deployment_blocked", + "_is_model_access_group_for_wildcard_route", + "_routing_groups", + "adaptive_routers", + "auto_routers", + "background_health_check_model_groups", + "cache", + "cached_deployment_model_info", + "cached_model_group_info", + "complexity_routers", + "config_deployments", + "cooldown_time", + "default_deployment", + "deployment_latency_map", + "deployment_names", + "deployments_for_request", + "enable_health_check_routing", + "fallback_access_check", + "fallback_budget_check", + "fallbacks", + "generate_model_id", + "get_available_guardrail", + "get_candidate_model_ids_for_route", + "get_configured_display_name", + "get_configured_mode", + "get_configured_service_tiers", + "get_credential_deployment", + "get_deployment", + "get_deployment_by_model_group_name", + "get_deployment_credentials", + "get_deployment_credentials_with_provider", + "get_deployment_model_info", + "get_discovered_model_info", + "get_fully_blocked_model_names", + "get_model_access_groups", + "get_model_group_info", + "get_model_ids", + "get_model_info", + "get_model_list", + "get_model_list_from_model_alias", + "get_model_listing_info", + "get_remaining_model_group_usage", + "get_routable_upstream_model", + "guardrail_list", + "has_model_id", + "health_check_ignore_transient_errors", + "health_state_cache", + "is_recognized_model", + "map_team_model", + "model_group_alias", + "model_name_to_deployment_indices", + "num_retries", + "pattern_router", + "provider_budget_config", + "quality_routers", + "resolve_model_name_from_model_id", + "resolved_litellm_models", + "retry_after", + "routable_model_group", + "router_general_settings", + "routing_plugins", + "search_tools", + "should_include_deployment", + "team_model_to_deployment_indices", + "team_pattern_routers", + "team_public_model_names", + "timeout", + } +) + +ASSIGNABLE: Final = frozenset({"cache_responses", "guardrail_list", "search_tools"}) diff --git a/tests/test_litellm_rust/router/test_proxy_switchover.py b/tests/test_litellm_rust/router/test_proxy_switchover.py index f74a8806b7c..49f49f9416b 100644 --- a/tests/test_litellm_rust/router/test_proxy_switchover.py +++ b/tests/test_litellm_rust/router/test_proxy_switchover.py @@ -8,6 +8,12 @@ from typing import Final import pytest +from litellm.router import Router +from litellm.router_backends.python_router import PythonRouter +from litellm.router_backends.rust_router import RustRouter, RustRouterUnsupportedError +from litellm.router_backends.selection import pinned_backend +from litellm.rust_bridge.configuration import Decision +from litellm.types.router import Deployment, RouterGeneralSettings from tests.test_litellm_rust.support.router_backends import BACKENDS, Backend, router_headers pytestmark = pytest.mark.requires_rust_extension @@ -25,6 +31,10 @@ def _deployment( } +def _content(response: object) -> str: + return str(response.choices[0].message.content) # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType,reportUnknownArgumentType] # ModelResponse + + class _Check: """A fallback access or budget check that records what it was asked and rejects `denied`.""" @@ -109,3 +119,117 @@ def test_the_sync_path_runs_the_checks_too(build: Callable[[Mapping[str, object] assert response.choices[0].message.content == "from c" # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType] # ModelResponse assert [model for model, _, _ in access.asked] == ["h", "i"] assert [model for model, _, _ in budget.asked] == ["i"] + + +def _facade(arguments: Mapping[str, object], decision: Decision = Decision.RUST_WITH_FALLBACK) -> Router: + with pinned_backend(decision): + return Router(**arguments) + + +def _plain(arguments: Mapping[str, object]) -> Mapping[str, object]: + return {key: list(value) if isinstance(value, tuple) else value for key, value in arguments.items()} + + +_TWO_GROUPS: Final = _plain({"model_list": (_deployment("g", "a"), _deployment("g", "b"), _deployment("h", "c"))}) +_READS: Final = ( + ("get_model_list", lambda router: router.get_model_list()), + ("get_model_names", lambda router: router.get_model_names()), + ("get_model_ids", lambda router: router.get_model_ids()), + ("get_deployment", lambda router: router.get_deployment(model_id="b")), + ("get_model_group_info", lambda router: router.get_model_group_info(model_group="g")), + ("get_model_access_groups", lambda router: router.get_model_access_groups()), + ("deployment_names", lambda router: router.deployment_names), + ("is_recognized_model", lambda router: router.is_recognized_model("h")), + ("get_fully_blocked_model_names", lambda router: router.get_fully_blocked_model_names()), +) + + +@pytest.mark.parametrize(("name", "read"), _READS, ids=[name for name, _ in _READS]) +async def test_the_proxys_reads_answer_as_the_python_router_does_and_stay_on_rust( + name: str, read: Callable[[Router], object] +) -> None: + del name + rust: Final = _facade(_TWO_GROUPS) + python: Final = _facade(_TWO_GROUPS, Decision.PYTHON) + for router in (rust, python): + router.upsert_deployment(Deployment(**_deployment("h", "d"))) + + assert read(rust) == read(python) + assert isinstance(rust.backend, RustRouter) + + +async def test_an_operation_rust_does_not_serve_hands_the_instance_over_for_good() -> None: + router: Final = _facade(_TWO_GROUPS) + router.search_tools = [{"search_tool_name": "web", "litellm_params": {"search_provider": "tavily"}}] + router.update_settings(num_retries=4) + rust: Final = router.backend + + response: Final = await router.aembedding(model="g", input=["hi"], mock_response=[0.5]) + + assert response.data[0]["embedding"] == [0.5] # pyright: ignore[reportUnknownMemberType,reportAttributeAccessIssue] # EmbeddingResponse + assert isinstance(router.backend, PythonRouter) + assert (router.search_tools, router.num_retries) == (rust.search_tools, 4) # pyright: ignore[reportAttributeAccessIssue] # Router views + assert _content(await router.acompletion("h", MESSAGES)) == "from c" + assert getattr(rust, "num_retries") == 4 + + +@pytest.mark.parametrize( + "option", + ( + pytest.param({"specific_deployment": True, "model": "openai/c"}, id="specific-deployment"), + pytest.param({"api_key": "client-key", "api_base": "https://example.invalid"}, id="client-side-credentials"), + pytest.param({"routing_strategy": "least-busy"}, id="per-request-routing-strategy"), + ), +) +async def test_a_request_option_rust_does_not_serve_is_answered_by_the_python_router( + option: Mapping[str, object], +) -> None: + router: Final = _facade(_TWO_GROUPS) + model: Final = str(option.get("model", "h")) + + response: Final = await router.acompletion( + model, MESSAGES, **{key: value for key, value in option.items() if key != "model"} + ) + + assert _content(response) == "from c" + assert isinstance(router.backend, PythonRouter) + + +def test_required_mode_raises_instead_of_handing_over() -> None: + router: Final = _facade(_TWO_GROUPS, Decision.RUST_REQUIRED) + + with pytest.raises(RustRouterUnsupportedError, match=r"Router\.aembedding"): + _ = router.aembedding + + assert isinstance(router.backend, RustRouter) + + +def test_an_assignable_setting_stays_on_rust() -> None: + router: Final = _facade(_TWO_GROUPS) + + router.cache_responses = True + + assert isinstance(router.backend, RustRouter) + assert router.backend._normalizer.cache_responses is True # pyright: ignore[reportPrivateUsage] # where attempts read it + + +def test_the_proxys_constructor_arguments_stay_on_rust_and_read_like_python() -> None: + access: Final = _Check(frozenset()) + arguments: Final = { + "model_list": [ + *_TWO_GROUPS["model_list"], # pyright: ignore[reportGeneralTypeIssues] # a list + {"model_name": "broken", "litellm_params": {"model": "no-such-provider-model"}}, + ], + "router_general_settings": RouterGeneralSettings(async_only_mode=True), + "search_tools": [], + "ignore_invalid_deployments": True, + "fallback_access_check": access, + "fallback_budget_check": access, + "auto_router_capability_limit": lambda: None, + } + rust: Final = _facade(arguments) + python: Final = _facade(arguments, Decision.PYTHON) + + assert isinstance(rust.backend, RustRouter) + assert (rust.model_names, rust.get_model_list()) == (python.model_names, python.get_model_list()) + assert rust.router_general_settings.async_only_mode is True diff --git a/tests/test_litellm_rust/router/test_redis_parity.py b/tests/test_litellm_rust/router/test_redis_parity.py index 8a094830bd0..2fc2e6cfca7 100644 --- a/tests/test_litellm_rust/router/test_redis_parity.py +++ b/tests/test_litellm_rust/router/test_redis_parity.py @@ -15,6 +15,7 @@ import pytest import redis import litellm +from litellm.caching.redis_cache import RedisCache from litellm.router_backends.python_router import PythonRouter from litellm.router_backends.rust_router import NATIVE_ROUTER, RustRouter from litellm.types.router import RouterRateLimitError @@ -147,3 +148,30 @@ async def test_a_rust_router_honors_cooldowns_a_python_router_wrote(store: redis rust_outcomes: Final = await _run(RustRouter(_arguments(redis_port, responses, {}), native, 0), 1) assert rust_outcomes == ["rejected"] + + +async def test_a_redis_attached_after_construction_carries_the_rust_routers_cooldowns( + store: redis.Redis, redis_port: int +) -> None: + """The proxy builds its router without Redis and attaches its own (`_update_redis_cache`) at startup.""" + responses: Final = ("litellm.RateLimitError", "litellm.RateLimitError") + native: Final = NATIVE_ROUTER.load() + assert native is not None + without_redis: Final = { + key: value for key, value in _arguments(redis_port, responses, {}).items() if not key.startswith("redis_") + } + rust_router: Final = RustRouter(without_redis, native, 0) + + rust_router._update_redis_cache(RedisCache(host="127.0.0.1", port=redis_port)) # pyright: ignore[reportPrivateUsage] # the proxy's own hand-off + rust_outcomes: Final = await _run(rust_router, 2) + python_router: Final = PythonRouter(**_arguments(redis_port, responses, {})) + python_outcomes: Final = await _run(python_router, 1) + python_router.discard() + + assert rust_outcomes == ["rate_limited", "rate_limited"] + assert sorted(key for key in store.keys("deployment:*:cooldown")) == [ + "deployment:d0:cooldown", + "deployment:d1:cooldown", + ] + assert python_outcomes == ["rejected"] + assert rust_router.cache.redis_cache is not None # pyright: ignore[reportAttributeAccessIssue] # a Router view diff --git a/tests/test_litellm_rust/router/test_rust_router.py b/tests/test_litellm_rust/router/test_rust_router.py index c31efda86d0..ebca5da9c86 100644 --- a/tests/test_litellm_rust/router/test_rust_router.py +++ b/tests/test_litellm_rust/router/test_rust_router.py @@ -11,7 +11,7 @@ from unittest.mock import MagicMock, patch import pytest import litellm -from litellm.router_backends.rust_router import RustRouter +from litellm.router_backends.rust_router import RustRouter, RustRouterUnsupportedError from litellm.router_backends.selection import select_backend from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.configuration import Rollout @@ -172,7 +172,7 @@ async def test_the_facade_serves_a_supported_config_from_the_rust_backend() -> N assert isinstance(backend, RustRouter) assert _content(await backend.acompletion("g", MESSAGES)) == "from a" assert backend.get_model_names() == ["g"] - with pytest.raises(NotImplementedError, match=r"Router\.cooldown_cache"): + with pytest.raises(RustRouterUnsupportedError, match=r"Router\.cooldown_cache"): _ = backend.cooldown_cache assert not hasattr(backend, "not_a_router_member")