diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 4175479f068..ba7fac47ba5 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -46,6 +46,7 @@ from litellm.secret_managers.main import get_secret, get_secret_str from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams if TYPE_CHECKING: + from litellm.caching.in_memory_cache import InMemoryCache from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.llms.openai import AllMessageValues @@ -1679,9 +1680,10 @@ class BedrockModelInfo(BaseLLMModelInfo): global_config = AmazonBedrockGlobalConfig() all_global_regions = global_config.get_all_regions() - def __init__(self, client: HTTPHandler | None = None) -> None: + def __init__(self, client: HTTPHandler | None = None, endpoint_failures: InMemoryCache | None = None) -> None: super().__init__() self._client: Final = client + self._endpoint_failures: Final = endpoint_failures @staticmethod def get_api_base(api_base: str | None = None) -> str | None: @@ -1715,10 +1717,17 @@ class BedrockModelInfo(BaseLLMModelInfo): def discover_models( self, litellm_params: Mapping[str, object] | None = None ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override - from litellm.llms.bedrock.model_listing import BedrockModelLister + from litellm.llms.bedrock.model_listing import CONTROL_PLANE_ENDPOINT_FAILURES, BedrockModelLister client: Final = self._client if self._client is not None else litellm.module_level_client - return sorted(BedrockModelLister(deployment=litellm_params or {}, client=client).invocable_model_ids()) + endpoint_failures: Final = ( + self._endpoint_failures if self._endpoint_failures is not None else CONTROL_PLANE_ENDPOINT_FAILURES + ) + return sorted( + BedrockModelLister( + deployment=litellm_params or {}, client=client, endpoint_failures=endpoint_failures + ).invocable_model_ids() + ) # def get_provider_info(self, model: str) -> Optional[ProviderSpecificModelInfo]: # """ diff --git a/litellm/llms/bedrock/model_listing.py b/litellm/llms/bedrock/model_listing.py index 8a35dff74ce..4ab4e4f4c2d 100644 --- a/litellm/llms/bedrock/model_listing.py +++ b/litellm/llms/bedrock/model_listing.py @@ -1,12 +1,15 @@ +import hashlib +import json from collections.abc import Iterator, Mapping from types import MappingProxyType -from typing import Final, TypeAlias +from typing import Final, Protocol, TypeAlias from urllib.parse import quote, urlencode import httpx from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from litellm._logging import verbose_logger +from litellm.caching.in_memory_cache import InMemoryCache from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token from litellm.llms.bedrock.common_utils import BedrockError, bedrock_control_plane_base_url from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -15,6 +18,8 @@ from litellm.types.llms.bedrock import AwsAuthParams LIST_MODELS_TIMEOUT: Final = 10.0 INFERENCE_PROFILES_PAGE_SIZE: Final = 1000 INFERENCE_PROFILES_PAGE_CAP: Final = 20 +CONTROL_PLANE_ENDPOINT_RETRY_SECONDS: Final = 3600 +CONTROL_PLANE_ENDPOINT_FAILURES: Final = InMemoryCache(default_ttl=CONTROL_PLANE_ENDPOINT_RETRY_SECONDS) _DECODED_JSON: Final = TypeAdapter(object) QueryPairs: TypeAlias = tuple[tuple[str, str], ...] @@ -46,8 +51,18 @@ class ListInferenceProfilesResponse(BaseModel): next_token: str | None = Field(default=None, alias="nextToken") +class EndpointFailureMemory(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: object) -> None: ... + + def delete_cache(self, key: str) -> None: ... + + class BedrockModelLister(BaseAWSLLM): - def __init__(self, deployment: Mapping[str, object], client: HTTPHandler) -> None: + def __init__( + self, deployment: Mapping[str, object], client: HTTPHandler, endpoint_failures: EndpointFailureMemory + ) -> None: super().__init__() self._aws_params: Final = dict(deployment) self._aws_region_name: Final = self._get_aws_region_name(optional_params=self._aws_params) @@ -57,27 +72,61 @@ class BedrockModelLister(BaseAWSLLM): configured_base_url: Final = bedrock_control_plane_base_url( deployment.get("aws_bedrock_control_plane_endpoint"), self._aws_region_name ) - self._sigv4_endpoint_base_url: Final = ( - None - if self._bearer_token is not None or configured_base_url == self._regional_base_url - else configured_base_url - ) + own_endpoint: Final = None if configured_base_url == self._regional_base_url else configured_base_url + self._sigv4_endpoint_base_url: Final = None if self._bearer_token is not None else own_endpoint + self._endpoint_skipped_for_api_key: Final = own_endpoint if self._bearer_token is not None else None self._client: Final = client + self._endpoint_failures: Final = endpoint_failures + self._deployment_key: Final = hashlib.sha256( + json.dumps(self._aws_params, sort_keys=True, default=str).encode() + ).hexdigest() def invocable_model_ids(self) -> frozenset[str]: - if self._sigv4_endpoint_base_url is None: + endpoint: Final = self._sigv4_endpoint_base_url + if endpoint is None: + return self._regional_model_ids() + if self._endpoint_failures.get_cache(self._deployment_key) is None: + return self._list_through(endpoint, self._regional_base_url) + return self._list_through(self._regional_base_url, endpoint) + + def _regional_model_ids(self) -> frozenset[str]: + skipped: Final = self._endpoint_skipped_for_api_key + if skipped is None: return self._invocable_model_ids(self._regional_base_url) try: - return self._invocable_model_ids(self._sigv4_endpoint_base_url) - except (BedrockError, ValueError) as error: - verbose_logger.warning( - "Listing Bedrock models through the configured control-plane endpoint %s failed, " - "so listing through %s instead: %s", - httpx.URL(self._sigv4_endpoint_base_url).host, - httpx.URL(self._regional_base_url).host, - error, - ) return self._invocable_model_ids(self._regional_base_url) + except BedrockError as error: + raise BedrockError( + status_code=error.status_code, + message=( + f"{error.message} A deployment with a Bedrock API key lists through " + f"{httpx.URL(self._regional_base_url).host}, never through the configured control-plane " + f"endpoint {httpx.URL(skipped).host}." + ), + ) from error + + def _list_through(self, first_base_url: str, fallback_base_url: str) -> frozenset[str]: + try: + listed: Final = self._invocable_model_ids(first_base_url) + except (BedrockError, ValueError) as error: + return self._list_through_fallback(first_base_url, fallback_base_url, error) + if listed: + return listed + return self._list_through_fallback(first_base_url, fallback_base_url, "it listed no models") + + def _list_through_fallback(self, failed_base_url: str, fallback_base_url: str, reason: object) -> frozenset[str]: + verbose_logger.warning( + "Listing Bedrock models through %s failed, so listing through %s instead: %s", + httpx.URL(failed_base_url).host, + httpx.URL(fallback_base_url).host, + reason, + ) + listed: Final = self._invocable_model_ids(fallback_base_url) + if failed_base_url == self._regional_base_url: + self._endpoint_failures.delete_cache(self._deployment_key) + else: + self._endpoint_failures.set_cache(self._deployment_key, True) + return listed def _invocable_model_ids(self, base_url: str) -> frozenset[str]: return frozenset(self._active_inference_profile_ids(base_url)) | frozenset( diff --git a/tests/integration/providers/test_bedrock_model_discovery_wire.py b/tests/integration/providers/test_bedrock_model_discovery_wire.py index 280f7638510..4d01dba552d 100644 --- a/tests/integration/providers/test_bedrock_model_discovery_wire.py +++ b/tests/integration/providers/test_bedrock_model_discovery_wire.py @@ -158,7 +158,7 @@ def test_admin_models_list_through_the_configured_control_plane_endpoint_only(ri assert_sigv4(requests, key=key, secret=secret, region="us-east-1", host=CONFIGURED_CONTROL_PLANE_HOST) -@pytest.mark.parametrize("status", (403, 404, 500)) +@pytest.mark.parametrize("status", (200, 403, 404, 500)) def test_configured_endpoint_failure_lists_through_the_regional_host(rig: Rig, status: int) -> None: key, secret = credential() marker: Final = stem() diff --git a/tests/unit/llms/bedrock/test_bedrock_model_listing.py b/tests/unit/llms/bedrock/test_bedrock_model_listing.py index 27159605e20..e78d3f9f9e7 100644 --- a/tests/unit/llms/bedrock/test_bedrock_model_listing.py +++ b/tests/unit/llms/bedrock/test_bedrock_model_listing.py @@ -10,7 +10,9 @@ from urllib.parse import parse_qsl, quote import httpx import pytest +from litellm.caching.in_memory_cache import InMemoryCache from litellm.llms.bedrock.common_utils import BedrockError, BedrockModelInfo +from litellm.llms.bedrock.model_listing import CONTROL_PLANE_ENDPOINT_RETRY_SECONDS from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.router import LiteLLM_Params @@ -130,8 +132,28 @@ def _hosts(handler: Mock) -> tuple[str, ...]: return tuple(call.args[0].url.host for call in handler.call_args_list) -def _bedrock(handler: Callable[[httpx.Request], httpx.Response]) -> BedrockModelInfo: - return BedrockModelInfo(client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler)))) +def _bedrock( + handler: Callable[[httpx.Request], httpx.Response], endpoint_failures: InMemoryCache | None = None +) -> BedrockModelInfo: + return BedrockModelInfo( + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + endpoint_failures=endpoint_failures if endpoint_failures is not None else InMemoryCache(), + ) + + +def _endpoint_answering( + endpoint_answer: Callable[[httpx.Request], httpx.Response], +) -> Callable[[httpx.Request], httpx.Response]: + regional: Final = _control_plane() + + def handler(request: httpx.Request) -> httpx.Response: + return endpoint_answer(request) if request.url.host == VPC_ENDPOINT_HOST else regional(request) + + return handler + + +def _forbidden(request: httpx.Request) -> httpx.Response: + return httpx.Response(403, json={"message": "Missing Authentication Token"}) def _deployment( @@ -211,19 +233,16 @@ def test_configured_endpoint_without_the_list_routes_falls_back_to_the_regional_ lambda request: httpx.Response(429, json={"message": "Too many requests"}), lambda request: httpx.Response(503, json={"message": "Service unavailable"}), lambda request: httpx.Response(200, text="batch gateway"), + lambda request: httpx.Response(200, json={"message": "ok"}), + lambda request: httpx.Response(200, json={"modelSummaries": [], "inferenceProfileSummaries": []}), _refuse_connection, ], - ids=["401", "403", "429", "503", "not-json", "connection-refused"], + ids=["401", "403", "429", "503", "not-json", "json-without-models", "empty-lists", "connection-refused"], ) def test_any_failure_through_the_configured_endpoint_lists_through_the_regional_host( endpoint_answer: Callable[[httpx.Request], httpx.Response], caplog: pytest.LogCaptureFixture ) -> None: - regional: Final = _control_plane() - - def endpoint_fails(request: httpx.Request) -> httpx.Response: - return endpoint_answer(request) if request.url.host == VPC_ENDPOINT_HOST else regional(request) - - handler: Final = Mock(wraps=endpoint_fails) + handler: Final = Mock(wraps=_endpoint_answering(endpoint_answer)) with caplog.at_level("WARNING", logger="LiteLLM"): assert ( @@ -237,6 +256,63 @@ def test_any_failure_through_the_configured_endpoint_lists_through_the_regional_ ) +def test_failed_configured_endpoint_is_skipped_until_the_retry_window_ends() -> None: + clock: Final = Mock(return_value=0.0) + handler: Final = Mock(wraps=_endpoint_answering(_forbidden)) + bedrock: Final = _bedrock(handler, InMemoryCache(default_ttl=CONTROL_PLANE_ENDPOINT_RETRY_SECONDS, clock=clock)) + deployment: Final = _deployment(control_plane_endpoint=f"https://{VPC_ENDPOINT_HOST}") + + assert bedrock.discover_models(deployment) == EXPECTED_MODELS + handler.reset_mock() + clock.return_value = CONTROL_PLANE_ENDPOINT_RETRY_SECONDS - 1.0 + assert bedrock.discover_models(deployment) == EXPECTED_MODELS + assert set(_hosts(handler)) == {REGIONAL_HOST} + + handler.reset_mock() + clock.return_value = CONTROL_PLANE_ENDPOINT_RETRY_SECONDS + 1.0 + assert bedrock.discover_models(deployment) == EXPECTED_MODELS + assert _hosts(handler)[0] == VPC_ENDPOINT_HOST + + +def test_endpoint_failure_of_one_deployment_leaves_another_on_its_endpoint() -> None: + endpoint_failures: Final = InMemoryCache() + deployment: Final = _deployment(control_plane_endpoint=f"https://{VPC_ENDPOINT_HOST}") + _bedrock(_endpoint_answering(_forbidden), endpoint_failures).discover_models(deployment) + endpoint_only: Final = Mock(wraps=_control_plane(host=VPC_ENDPOINT_HOST)) + + assert ( + _bedrock(endpoint_only, endpoint_failures).discover_models({**deployment, "model": "bedrock/anthropic.*"}) + == EXPECTED_MODELS + ) + assert set(_hosts(endpoint_only)) == {VPC_ENDPOINT_HOST} + + +def test_regional_failure_while_the_endpoint_is_skipped_returns_listing_to_the_endpoint() -> None: + endpoint_failures: Final = InMemoryCache() + deployment: Final = _deployment(control_plane_endpoint=f"https://{VPC_ENDPOINT_HOST}") + _bedrock(_endpoint_answering(_forbidden), endpoint_failures).discover_models(deployment) + endpoint_only: Final = Mock(wraps=_control_plane(host=VPC_ENDPOINT_HOST)) + bedrock: Final = _bedrock(endpoint_only, endpoint_failures) + + assert bedrock.discover_models(deployment) == EXPECTED_MODELS + assert _hosts(endpoint_only)[0] == REGIONAL_HOST + endpoint_only.reset_mock() + assert bedrock.discover_models(deployment) == EXPECTED_MODELS + assert set(_hosts(endpoint_only)) == {VPC_ENDPOINT_HOST} + + +@pytest.mark.parametrize("control_plane_endpoint", [None, f"https://{VPC_ENDPOINT_HOST}"], ids=["regional", "configured"]) +def test_api_key_listing_failure_names_only_a_configured_endpoint_it_skipped( + control_plane_endpoint: str | None, +) -> None: + with pytest.raises(BedrockError, match=rf"{re.escape(REGIONAL_HOST)} .*ConnectError") as raised: + _bedrock(_refuse_connection).discover_models( + _deployment(api_key=BEARER_TOKEN, control_plane_endpoint=control_plane_endpoint) + ) + + assert (VPC_ENDPOINT_HOST in str(raised.value)) == (control_plane_endpoint is not None) + + def test_bearer_token_api_key_replaces_sigv4() -> None: bedrock: Final = _bedrock(_control_plane(bearer_token=BEARER_TOKEN))