diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 6ca02cbbe64..657fd3abfbc 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -27,6 +27,7 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( "aws_external_id", "aws_session_tags", "aws_bedrock_runtime_endpoint", + "aws_bedrock_control_plane_endpoint", "aws_bedrock_project_id", } ) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 81057a1ee23..9783ac2aa10 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -5,7 +5,10 @@ from typing import TYPE_CHECKING, Final, Literal from openai.types.batch import BatchRequestCounts from openai.types.batch import Metadata as OpenAIBatchMetadata -from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.llms.bedrock.common_utils import ( + bedrock_control_plane_base_url, + configured_bedrock_control_plane_endpoint, +) from litellm.types.llms.bedrock import AwsAuthParams, AwsSessionTag from litellm.types.utils import LiteLLMBatch @@ -156,10 +159,14 @@ class BedrockBatchesHandler: aws_session_tags=aws_session_tags, ) creds: Final = BedrockBatchesConfig().resolve_credentials(auth_params, region) + control_plane_endpoint: Final = configured_bedrock_control_plane_endpoint( + kwargs.get("aws_bedrock_control_plane_endpoint") + ) client: Final = boto3.client( "bedrock", region_name=region, + endpoint_url=control_plane_endpoint, aws_access_key_id=creds.access_key, aws_secret_access_key=creds.secret_key, aws_session_token=creds.token, @@ -171,6 +178,7 @@ class BedrockBatchesHandler: batch_id=batch_id, aws_region_name=region, logging_obj=logging_obj, + aws_bedrock_control_plane_endpoint=control_plane_endpoint, **auth_params.model_dump(), ) @@ -270,6 +278,7 @@ class BedrockBatchesHandler: batch_id: str, aws_region_name: str | None = None, logging_obj=None, + aws_bedrock_control_plane_endpoint: str | None = None, **kwargs, ) -> "LiteLLMBatch": """ @@ -289,6 +298,8 @@ class BedrockBatchesHandler: aws_region_name: Region for the boto3 ``bedrock`` client. If omitted, we fall back to parsing the region out of ``batch_id`` itself. logging_obj: Optional litellm logging object. + aws_bedrock_control_plane_endpoint: Sends the call to a VPC endpoint + or gateway instead of the regional ``bedrock`` host. **kwargs: Optional AWS credential overrides (``aws_access_key_id``, ``aws_secret_access_key``, ``aws_session_token``, ``aws_profile_name``, @@ -322,6 +333,7 @@ class BedrockBatchesHandler: client: Final = boto3.client( "bedrock", region_name=region, + endpoint_url=configured_bedrock_control_plane_endpoint(aws_bedrock_control_plane_endpoint), aws_access_key_id=creds.access_key, aws_secret_access_key=creds.secret_key, aws_session_token=creds.token, @@ -339,9 +351,7 @@ class BedrockBatchesHandler: api_key="", additional_args={ "complete_input_dict": {"jobIdentifier": batch_id}, - "api_base": ( - f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{url_path_id}" - ), + "api_base": f"{bedrock_control_plane_base_url(aws_bedrock_control_plane_endpoint, region)}/model-invocation-job/{url_path_id}", }, ) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 3ad8049ec26..689db5f7ca2 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Final, Literal, cast from httpx import Headers, Response from pydantic import TypeAdapter, ValidationError -from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix, is_bedrock_arn +from litellm.litellm_core_utils.aws_partition import is_bedrock_arn from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, ) @@ -32,6 +32,7 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders, Usage from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( CommonBatchFilesUtils, + bedrock_control_plane_base_url, merge_bedrock_aws_request_params, resolve_s3_bucket_owner, resolve_s3_encryption_key_id, @@ -178,10 +179,8 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): aws_region_name: Final = self._get_aws_region_name(request_params, model) # Bedrock model invocation job endpoint - # Format: https://bedrock.{region}.{partition dns suffix}/model-invocation-job - bedrock_endpoint: Final = ( - f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job" - ) + # Format: {control plane base}/model-invocation-job + bedrock_endpoint: Final = f"{bedrock_control_plane_base_url(request_params.get('aws_bedrock_control_plane_endpoint'), aws_region_name)}/model-invocation-job" return bedrock_endpoint @@ -279,9 +278,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # Use common utility for AWS signing request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) aws_region_name: Final = self._get_aws_region_name(request_params, model) - endpoint_url: Final = ( - f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job" - ) + endpoint_url: Final = f"{bedrock_control_plane_base_url(request_params.get('aws_bedrock_control_plane_endpoint'), aws_region_name)}/model-invocation-job" signed_headers, signed_data = self.common_utils.sign_aws_request( service_name="bedrock", data=bedrock_request, @@ -431,12 +428,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): import urllib.parse as _ul encoded_arn: Final = _ul.quote(batch_id, safe="") - endpoint_url: Final = ( - f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{encoded_arn}" - ) + request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) + endpoint_url: Final = f"{bedrock_control_plane_base_url(request_params.get('aws_bedrock_control_plane_endpoint'), region)}/model-invocation-job/{encoded_arn}" # Use common utility for AWS signing - request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) signed_headers, _ = self.common_utils.sign_aws_request( service_name="bedrock", data={}, # GET request has no body diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index bbbc929fad3..0cf3e895ad2 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -138,6 +138,25 @@ def merge_bedrock_aws_request_params( return request_params +def configured_bedrock_control_plane_endpoint(endpoint: object) -> str | None: + """An ``aws_bedrock_control_plane_endpoint`` value without a trailing slash, or None when unset.""" + if isinstance(endpoint, str) and endpoint.strip(): + return endpoint.strip().rstrip("/") + return None + + +def bedrock_control_plane_base_url(endpoint: object, aws_region_name: str) -> str: + """Base URL for Bedrock control-plane calls such as model invocation jobs. + + A configured ``aws_bedrock_control_plane_endpoint`` (a VPC endpoint or a compatible gateway) + wins over the regional ``bedrock`` host. ``api_base`` is deliberately not used here, since on a + Bedrock deployment it points at the runtime host, which does not serve these APIs. + """ + return configured_bedrock_control_plane_endpoint(endpoint) or ( + f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" + ) + + def s3_static_key_pair(params: Mapping[str, object]) -> tuple[str, str] | None: """The s3_access_key_id / s3_secret_access_key pair when both are set, otherwise None.""" s3_access_key_id: Final = params.get("s3_access_key_id") diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 7a05121ff2b..cad156827d7 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -417,6 +417,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( # tokens) to the attacker's host, or coerces the proxy into # authenticating against the attacker's host with admin secrets. "aws_bedrock_runtime_endpoint", + "aws_bedrock_control_plane_endpoint", # Bedrock project/workspace association. Deployments pin this to # enforce a data-retention policy, so a caller-supplied value would # re-route the request's retention and accounting to any project diff --git a/litellm/types/router.py b/litellm/types/router.py index 36d298b2bc2..8ca6dd5cfc2 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -411,6 +411,7 @@ class CredentialLiteLLMParams(LiteLLMBaseModel): aws_external_id: str | None = None aws_session_tags: Sequence[AwsSessionTag] | None = None aws_bedrock_runtime_endpoint: str | None = None + aws_bedrock_control_plane_endpoint: str | None = None aws_bedrock_project_id: str | None = None s3_bucket_name: str | None = None s3_endpoint_url: str | None = None diff --git a/tests/unit/llms/bedrock/batches/test_handler.py b/tests/unit/llms/bedrock/batches/test_handler.py index d920eabcdf6..9a88350b5e8 100644 --- a/tests/unit/llms/bedrock/batches/test_handler.py +++ b/tests/unit/llms/bedrock/batches/test_handler.py @@ -257,6 +257,23 @@ def test_explicit_region_overrides_arn(patched_boto3): assert kwargs["region_name"] == "eu-central-1" +def test_status_client_uses_regional_host_when_no_control_plane_endpoint(patched_boto3): + _, boto_client_factory = patched_boto3 + BedrockBatchesHandler.handle_model_invocation_job_status(batch_id=JOB_ARN) + _, kwargs = boto_client_factory.call_args + assert kwargs["endpoint_url"] is None + + +def test_status_client_uses_configured_control_plane_endpoint(patched_boto3): + _, boto_client_factory = patched_boto3 + BedrockBatchesHandler.handle_model_invocation_job_status( + batch_id=JOB_ARN, + aws_bedrock_control_plane_endpoint="https://bedrock-gateway.internal/", + ) + _, kwargs = boto_client_factory.call_args + assert kwargs["endpoint_url"] == "https://bedrock-gateway.internal" + + def test_failure_message_propagates(patched_boto3): fake_client, _ = patched_boto3 failed_response = _fake_boto3_response(status="Failed") @@ -479,6 +496,28 @@ def test_litellm_cancel_batch_dispatches_to_bedrock(patched_boto3): assert batch.status == "cancelled" +def test_litellm_batch_calls_send_every_bedrock_client_to_the_configured_endpoint(patched_boto3): + import litellm + + fake_client, boto_client_factory = patched_boto3 + fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped") + + litellm.cancel_batch( + batch_id=JOB_ARN, + custom_llm_provider="bedrock", + aws_bedrock_control_plane_endpoint="https://bedrock-gateway.internal", + ) + litellm.retrieve_batch( + batch_id=JOB_ARN, + custom_llm_provider="bedrock", + aws_bedrock_control_plane_endpoint="https://bedrock-gateway.internal", + ) + + # cancel builds a client to stop the job and another to read its status, retrieve builds one more + endpoints = [call.kwargs["endpoint_url"] for call in boto_client_factory.call_args_list] + assert endpoints == ["https://bedrock-gateway.internal"] * 3 + + class _TagGatedSTSClient: """Stands in for STS behind a trust policy that only admits sessions carrying ``tags``.""" diff --git a/tests/unit/llms/bedrock/batches/test_transformation.py b/tests/unit/llms/bedrock/batches/test_transformation.py index 5e987239c54..b5efcfd975f 100644 --- a/tests/unit/llms/bedrock/batches/test_transformation.py +++ b/tests/unit/llms/bedrock/batches/test_transformation.py @@ -66,6 +66,89 @@ def test_get_complete_batch_url_uses_region(config): assert url == "https://bedrock.eu-central-1.amazonaws.com/model-invocation-job" +CONTROL_PLANE_ENDPOINT = "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/" + + +def test_get_complete_batch_url_uses_configured_control_plane_endpoint(config): + url = config.get_complete_batch_url( + api_base=None, + api_key=None, + model="anthropic.claude-3", + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={"aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT}, + data={"input_file_id": "s3://b/k"}, + ) + assert url == "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job" + + +def test_proxy_deployment_credentials_carry_the_control_plane_endpoint(config): + # the proxy hands batch calls the deployment's credentials through the strict CredentialLiteLLMParams filter + from litellm import Router + + router = Router( + model_list=[ + { + "model_name": "claude-batch", + "litellm_params": { + "model": "bedrock/anthropic.claude-3", + "aws_region_name": "us-west-2", + "aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT, + }, + } + ] + ) + credentials = router.get_deployment_credentials_with_provider("claude-batch") + + url = config.get_complete_batch_url( + api_base=None, + api_key=None, + model="anthropic.claude-3", + optional_params={}, + litellm_params=credentials, + data={"input_file_id": "s3://b/k"}, + ) + assert url == "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job" + + +def test_get_complete_batch_url_ignores_runtime_api_base(config): + # api_base on a Bedrock deployment points at the runtime host, which does not serve batch jobs + url = config.get_complete_batch_url( + api_base="https://bedrock-runtime.us-west-2.amazonaws.com", + api_key=None, + model="anthropic.claude-3", + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={"api_base": "https://bedrock-runtime.us-west-2.amazonaws.com"}, + data={"input_file_id": "s3://b/k"}, + ) + assert url == "https://bedrock.us-west-2.amazonaws.com/model-invocation-job" + + +def test_create_request_sends_and_signs_for_configured_control_plane_endpoint(config): + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-deadbeef", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({"Authorization": "signed"}, b"{}") + result = config.transform_create_batch_request( + model="anthropic.claude-3-5-sonnet", + create_batch_data={ + "input_file_id": "s3://in-bucket/input.jsonl", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::123:role/my-batch-role", + "aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT, + }, + ) + + expected = "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job" + assert result["url"] == expected + assert mock_sign.call_args.kwargs["endpoint_url"] == expected + + # --------------------------------------------------------------------------- # # transform_create_batch_request - request construction (sign_aws_request mocked) # --------------------------------------------------------------------------- # @@ -608,6 +691,23 @@ def test_retrieve_request_builds_encoded_arn_url(config): assert mock_sign.call_args.kwargs["data"] == {} +def test_retrieve_request_sends_and_signs_for_configured_control_plane_endpoint(config): + with patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({"Authorization": "signed"}, b"") + result = config.transform_retrieve_batch_request( + batch_id=ARN, + optional_params={}, + litellm_params={"aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT}, + ) + + expected = ( + "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job/" + "arn%3Aaws%3Abedrock%3Aus-west-2%3A123456789012%3Amodel-invocation-job%2Fabc1234567" + ) + assert result["url"] == expected + assert mock_sign.call_args.kwargs["endpoint_url"] == expected + + def test_retrieve_request_rejects_non_arn(config): with pytest.raises(ValueError, match="Expected ARN"): config.transform_retrieve_batch_request( diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 02be27d921e..f037aa28abd 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -3357,6 +3357,7 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: "field", [ "aws_bedrock_runtime_endpoint", + "aws_bedrock_control_plane_endpoint", "langsmith_base_url", "langfuse_host", "posthog_host", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 733a093bcac..f6fefbf7d28 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35028,6 +35028,8 @@ export interface components { aws_access_key_id?: string | null; /** Aws Batch Role Arn */ aws_batch_role_arn?: string | null; + /** Aws Bedrock Control Plane Endpoint */ + aws_bedrock_control_plane_endpoint?: string | null; /** Aws Bedrock Project Id */ aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */ @@ -50278,6 +50280,8 @@ export interface components { aws_access_key_id?: string | null; /** Aws Batch Role Arn */ aws_batch_role_arn?: string | null; + /** Aws Bedrock Control Plane Endpoint */ + aws_bedrock_control_plane_endpoint?: string | null; /** Aws Bedrock Project Id */ aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */