From 21e06501e633c7c457841c336b8c2a79f4eef09d Mon Sep 17 00:00:00 2001 From: Priyansh Nandwana Date: Sun, 18 Oct 2026 14:55:28 +0530 Subject: [PATCH] fix(bedrock): send batch jobs to a configured control-plane endpoint Bedrock batch create, retrieve and cancel always built the regional bedrock.{region} host, so a VPC endpoint or compatible gateway could not be used. Add aws_bedrock_control_plane_endpoint as a deployment setting read by the signed HTTP path and the boto3 path, carry it through the proxy credential filter, and refuse it in request bodies like the other endpoint overrides. api_base stays out of it because on a Bedrock deployment it points at the runtime host. Fixes #45056 --- .../litellm_core_utils/get_litellm_params.py | 1 + litellm/llms/bedrock/batches/handler.py | 18 +++- .../llms/bedrock/batches/transformation.py | 19 ++-- litellm/llms/bedrock/common_utils.py | 19 ++++ litellm/proxy/auth/auth_utils.py | 1 + litellm/types/router.py | 1 + .../unit/llms/bedrock/batches/test_handler.py | 39 +++++++ .../bedrock/batches/test_transformation.py | 100 ++++++++++++++++++ tests/unit/proxy/auth/test_auth_utils.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 10 files changed, 187 insertions(+), 16 deletions(-) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 9dee3714c60..be543496a14 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 94cf7ff4dbe..3aa416e6837 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -409,6 +409,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 1bd1c2e5502..3b37251b1f9 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -401,6 +401,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 7de249bfb6e..748a7b17009 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -3018,6 +3018,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 c33ef86af72..380f7ba3d64 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34986,6 +34986,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 */ @@ -50228,6 +50230,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 */