From 898569b48f29607fa58ad4292dab6a74ec263db9 Mon Sep 17 00:00:00 2001 From: Mubashir Osmani Date: Mon, 20 Jul 2026 19:55:17 +0000 Subject: [PATCH] fix(bedrock/batches): propagate model, IAM role and AWS creds through proxy batch routing Bedrock create_batch through the proxy needs three values that model-based routing was dropping or mis-resolving, so it failed after config resolution inside the Bedrock batch path. CredentialLiteLLMParams re-validates a deployment's litellm_params and silently drops any undeclared field. aws_batch_role_arn (plus session token, assumed-role, and S3 output/region/endpoint/KMS fields) were undeclared, so the resolved credentials lost the IAM role and Bedrock raised 'AWS IAM role ARN is required'. Declare the AWS/S3 batch fields so they survive resolution. base_llm_http_handler.create_batch calls the transform with optional_params={} and the deployment credentials in litellm_params, but the AWS signer reads only optional_params, so boto3 saw no keys ('Unable to locate credentials'). Merge litellm_params into the signing params (per-request optional_params still wins) for region resolution and signing. get_deployment_credentials_with_provider never returned the deployment's underlying model, so the batch endpoints forwarded the proxy model-group name (alias) as the Bedrock modelId, which AWS rejects. Return the deployment model; the endpoints already forward it. This matches what the endpoint tests already assert. Regression coverage: credential round-trip keeps the role and returns the model, and the transform signs with litellm_params credentials. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/bedrock/batches/transformation.py | 5 +- litellm/router.py | 2 + litellm/types/router.py | 11 ++ tests/e2e/batches/capabilities.py | 2 - .../bedrock/batches/test_transformation.py | 55 +++++++++ ...est_bedrock_batch_credential_resolution.py | 115 ++++++++++++++++++ 6 files changed, 186 insertions(+), 4 deletions(-) create mode 100644 tests/test_litellm/test_bedrock_batch_credential_resolution.py diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 4fcf7cf91cb..f53a3fd8cd5 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -211,14 +211,15 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # For Bedrock, we need to return a pre-signed request with AWS auth headers # Use common utility for AWS signing + signing_params = {**litellm_params, **optional_params} endpoint_url = ( - f"https://bedrock.{self._get_aws_region_name(optional_params, model)}.amazonaws.com/model-invocation-job" + f"https://bedrock.{self._get_aws_region_name(signing_params, model)}.amazonaws.com/model-invocation-job" ) signed_headers, signed_data = self.common_utils.sign_aws_request( service_name="bedrock", data=bedrock_request, endpoint_url=endpoint_url, - optional_params=optional_params, + optional_params=signing_params, method="POST", ) diff --git a/litellm/router.py b/litellm/router.py index ae3f7ba11c2..42e2ac309c5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8638,6 +8638,8 @@ class Router: else: credentials["custom_llm_provider"] = "openai" # default + credentials["model"] = deployment.litellm_params.model + return credentials @overload diff --git a/litellm/types/router.py b/litellm/types/router.py index 28e4a8272e8..f598ce9551d 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -207,10 +207,21 @@ class CredentialLiteLLMParams(BaseModel): ## AWS BEDROCK / SAGEMAKER ## aws_access_key_id: Optional[str] = None aws_secret_access_key: Optional[str] = None + aws_session_token: Optional[str] = None aws_region_name: Optional[str] = None aws_bedrock_runtime_endpoint: Optional[str] = None aws_bedrock_project_id: Optional[str] = None + aws_batch_role_arn: Optional[str] = None + aws_role_name: Optional[str] = None + aws_session_name: Optional[str] = None + aws_web_identity_token: Optional[str] = None + aws_sts_endpoint: Optional[str] = None + aws_profile_name: Optional[str] = None s3_bucket_name: Optional[str] = None + s3_output_bucket_name: Optional[str] = None + s3_region_name: Optional[str] = None + s3_endpoint_url: Optional[str] = None + s3_encryption_key_id: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index aeeb67b6169..3d8890bbfb7 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -75,8 +75,6 @@ class Provider: aws_region_name="os.environ/AWS_REGION", s3_region_name="os.environ/AWS_REGION", s3_bucket_name=_env_ref("AWS_BATCH_S3_BUCKET", "AWS_S3_BUCKET_NAME"), - s3_access_key_id="os.environ/AWS_ACCESS_KEY_ID", - s3_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", aws_batch_role_arn="os.environ/AWS_BATCH_ROLE_ARN", ) case _: diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index d1ad5943ae6..8e4f29e4a68 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -239,6 +239,61 @@ def test_create_request_missing_model_raises(config): ) +def test_create_request_signs_with_credentials_from_litellm_params(config): + """Regression: the proxy's model-based routing forwards deployment AWS + credentials via ``litellm_params`` while ``base_llm_http_handler.create_batch`` + calls this transform with ``optional_params={}``. The signer reads credentials + from ``optional_params``, so unless the transform merges ``litellm_params`` in, + boto3 sees no keys and raises "Unable to locate credentials"; region resolution + also has to see ``aws_region_name`` from ``litellm_params``.""" + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "aws_access_key_id": "AKIA-DEPLOYMENT", + "aws_secret_access_key": "secret-deployment", + "aws_region_name": "ap-south-1", + }, + ) + signing_params = mock_sign.call_args.kwargs["optional_params"] + assert signing_params["aws_access_key_id"] == "AKIA-DEPLOYMENT" + assert signing_params["aws_secret_access_key"] == "secret-deployment" + assert mock_sign.call_args.kwargs["endpoint_url"] == ( + "https://bedrock.ap-south-1.amazonaws.com/model-invocation-job" + ) + + +def test_create_request_optional_params_win_over_litellm_params_for_signing(config): + """Per-request ``optional_params`` must override deployment ``litellm_params`` + when both carry the same credential field.""" + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "aws_region_name": "ap-south-1", + }, + ) + assert mock_sign.call_args.kwargs["endpoint_url"] == ( + "https://bedrock.us-west-2.amazonaws.com/model-invocation-job" + ) + + def test_create_request_no_timeout_for_non_24h_window(config): with patch.object( config.common_utils, diff --git a/tests/test_litellm/test_bedrock_batch_credential_resolution.py b/tests/test_litellm/test_bedrock_batch_credential_resolution.py new file mode 100644 index 00000000000..52990b230f5 --- /dev/null +++ b/tests/test_litellm/test_bedrock_batch_credential_resolution.py @@ -0,0 +1,115 @@ +""" +Regression for Bedrock batch creation through the proxy. + +``Router.get_deployment_credentials_with_provider`` resolves the upstream +credentials for a deployment by model id for the proxy's ``/v1/files`` and +``/v1/batches`` routing. It builds the credentials by round-tripping +``litellm_params`` through ``CredentialLiteLLMParams``:: + + CredentialLiteLLMParams( + **deployment.litellm_params.model_dump(exclude_none=True) + ).model_dump(exclude_none=True) + +That re-validation is strict: any field not declared on +``CredentialLiteLLMParams`` is silently dropped. Two Bedrock-batch bugs came +from that plus a missing field on the resolved dict: + +1. ``aws_batch_role_arn`` (and the other AWS/S3 job fields) were undeclared, so + a Bedrock deployment lost its IAM role on the way to ``create_batch`` and the + transform raised "AWS IAM role ARN is required for Bedrock batch jobs". +2. The resolved credentials never carried the deployment's real ``model``, so + the batch endpoints forwarded the proxy model-group name (alias) as the + Bedrock ``modelId``. Bedrock create needs the real model id (e.g. + ``bedrock/us.anthropic...``); the alias is rejected by AWS. + +The tests below pin both: the strict credential model keeps the AWS/S3 batch +fields, and the resolver returns the deployment's underlying model. +""" + + +class TestCredentialLiteLLMParamsBedrockFields: + def test_aws_batch_role_arn_round_trips_through_model_dump(self): + from litellm.types.router import CredentialLiteLLMParams + + params = CredentialLiteLLMParams( + aws_access_key_id="AKIA-xyz", + aws_secret_access_key="secret-xyz", + aws_region_name="us-east-1", + aws_batch_role_arn="arn:aws:iam::123:role/batch", + s3_bucket_name="my-bucket", + ) + dumped = params.model_dump(exclude_none=True) + assert dumped["aws_batch_role_arn"] == "arn:aws:iam::123:role/batch", ( + "aws_batch_role_arn dropped from CredentialLiteLLMParams.model_dump(); " + "the proxy's model-based batch routing loses the IAM role and Bedrock " + "create_batch fails with 'AWS IAM role ARN is required'" + ) + assert dumped["aws_access_key_id"] == "AKIA-xyz" + assert dumped["aws_secret_access_key"] == "secret-xyz" + assert dumped["s3_bucket_name"] == "my-bucket" + + def test_bedrock_fields_are_optional(self): + from litellm.types.router import CredentialLiteLLMParams + + params = CredentialLiteLLMParams(api_key="sk-static") + dumped = params.model_dump(exclude_none=True) + assert "aws_batch_role_arn" not in dumped + assert dumped["api_key"] == "sk-static" + + +class TestRouterBedrockBatchCredentialResolution: + def _bedrock_router(self, deployment_id: str): + from litellm import Router + + return Router( + model_list=[ + { + "model_name": "bedrock-batch-sonnet", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-sonnet-4-6", + "aws_access_key_id": "AKIA-deployment", + "aws_secret_access_key": "secret-deployment", + "aws_region_name": "us-east-1", + "s3_bucket_name": "my-bucket", + "aws_batch_role_arn": "arn:aws:iam::123:role/batch", + }, + "model_info": {"id": deployment_id}, + } + ] + ) + + def test_credentials_preserve_aws_batch_role_arn(self): + router = self._bedrock_router("bedrock-batch-deployment") + + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch-deployment" + ) + assert credentials is not None + assert credentials["custom_llm_provider"] == "bedrock" + assert credentials.get("aws_batch_role_arn") == "arn:aws:iam::123:role/batch", ( + "Router credential resolution dropped aws_batch_role_arn; Bedrock " + "create_batch cannot build the model-invocation-job request" + ) + assert credentials.get("aws_access_key_id") == "AKIA-deployment" + assert credentials.get("aws_secret_access_key") == "secret-deployment" + + def test_credentials_include_deployment_model(self): + """The resolved credentials must carry the deployment's real model so the + batch endpoints forward it (not the proxy alias) as the Bedrock modelId.""" + router = self._bedrock_router("bedrock-batch-deployment") + + by_id = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch-deployment" + ) + assert by_id is not None + assert by_id.get("model") == "bedrock/us.anthropic.claude-sonnet-4-6", ( + "Router credential resolution did not return the deployment model; the " + "batch endpoints would send the proxy model-group name as the Bedrock " + "modelId, which AWS rejects" + ) + + by_group = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch-sonnet" + ) + assert by_group is not None + assert by_group.get("model") == "bedrock/us.anthropic.claude-sonnet-4-6"