Merge pull request #42262 from BerriAI/litellm_bedrock_batch_s3_bucket_owner

* fix(bedrock): send s3BucketOwner on batch input and output data config

Resolve s3_bucket_owner from litellm_params, then optional_params, then
AWS_S3_BUCKET_OWNER and emit it on both S3 data configs so cross-account
batch buckets pass Bedrock ownership validation. Omitted when unset

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(bedrock): build batch output config with explicit returns

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/llms/bedrock/batches/transformation.py

@@ -33,6 +33,7 @@ from ..base_aws_llm import BaseAWSLLM
 from ..common_utils import (
     CommonBatchFilesUtils,
     merge_bedrock_aws_request_params,
+    resolve_s3_bucket_owner,
     resolve_s3_encryption_key_id,
 )
 
@@ -51,6 +52,26 @@ _S3_BATCH_FILE_UUID_SUFFIX_PATTERN: Final = re.compile(
 _BEDROCK_TAGS_ADAPTER: Final[TypeAdapter[list[BedrockTag]]] = TypeAdapter(list[BedrockTag])
 
 
+def _build_s3_input_config(s3_uri: str, s3_bucket_owner: str | None) -> BedrockS3InputDataConfig:
+    if s3_bucket_owner is None:
+        return BedrockS3InputDataConfig(s3Uri=s3_uri)
+    return BedrockS3InputDataConfig(s3Uri=s3_uri, s3BucketOwner=s3_bucket_owner)
+
+
+def _build_s3_output_config(
+    s3_uri: str, s3_bucket_owner: str | None, s3_encryption_key_id: str | None
+) -> BedrockS3OutputDataConfig:
+    if s3_bucket_owner is None:
+        if s3_encryption_key_id is None:
+            return BedrockS3OutputDataConfig(s3Uri=s3_uri)
+        return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3EncryptionKeyId=s3_encryption_key_id)
+    if s3_encryption_key_id is None:
+        return BedrockS3OutputDataConfig(s3Uri=s3_uri, s3BucketOwner=s3_bucket_owner)
+    return BedrockS3OutputDataConfig(
+        s3Uri=s3_uri, s3BucketOwner=s3_bucket_owner, s3EncryptionKeyId=s3_encryption_key_id
+    )
+
+
 def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]:
     try:
         return _BEDROCK_TAGS_ADAPTER.validate_python(raw_tags, strict=True)
@@ -214,25 +235,23 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
         job_name: Final = self.common_utils.generate_unique_job_name(model, prefix="litellm")
         output_key: Final = f"litellm-batch-outputs/{job_name}/"
 
-        # Build input data config
-        input_data_config: Final[BedrockInputDataConfig] = {
-            "s3InputDataConfig": BedrockS3InputDataConfig(s3Uri=f"s3://{input_bucket}/{input_key}")
-        }
-
-        # Build output data config
-        s3_output_config: Final[BedrockS3OutputDataConfig] = BedrockS3OutputDataConfig(
-            s3Uri=f"s3://{output_bucket}/{output_key}"
-        )
-
-        # Add optional KMS encryption key ID if provided
-        s3_encryption_key_id = resolve_s3_encryption_key_id(
+        s3_bucket_owner: Final = resolve_s3_bucket_owner(litellm_params=litellm_params, optional_params=optional_params)
+        s3_encryption_key_id: Final = resolve_s3_encryption_key_id(
             litellm_params=litellm_params,
             optional_params=optional_params,
         )
-        if s3_encryption_key_id:
-            s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id
-
-        output_data_config: Final[BedrockOutputDataConfig] = {"s3OutputDataConfig": s3_output_config}
+        input_data_config: Final[BedrockInputDataConfig] = {
+            "s3InputDataConfig": _build_s3_input_config(
+                s3_uri=f"s3://{input_bucket}/{input_key}", s3_bucket_owner=s3_bucket_owner
+            )
+        }
+        output_data_config: Final[BedrockOutputDataConfig] = {
+            "s3OutputDataConfig": _build_s3_output_config(
+                s3_uri=f"s3://{output_bucket}/{output_key}",
+                s3_bucket_owner=s3_bucket_owner,
+                s3_encryption_key_id=s3_encryption_key_id,
+            )
+        }
 
         # Create Bedrock batch request with proper typing
         bedrock_request: Final[BedrockCreateBatchRequest] = {

litellm/llms/bedrock/common_utils.py

@@ -1555,11 +1555,33 @@ def resolve_s3_encryption_key_id(
     Precedence: `s3_encryption_key_id` in litellm_params, then optional_params
     (client-side / request params), then the AWS_S3_ENCRYPTION_KEY_ID env var.
     """
+    return _resolve_s3_setting("s3_encryption_key_id", "AWS_S3_ENCRYPTION_KEY_ID", litellm_params, optional_params)
+
+
+def resolve_s3_bucket_owner(
+    litellm_params: Mapping[str, object],
+    optional_params: Mapping[str, object] | None = None,
+) -> str | None:
+    """
+    Resolve the AWS account id that owns the S3 buckets used by Bedrock batch jobs.
+
+    Precedence: `s3_bucket_owner` in litellm_params, then optional_params
+    (client-side / request params), then the AWS_S3_BUCKET_OWNER env var.
+    """
+    return _resolve_s3_setting("s3_bucket_owner", "AWS_S3_BUCKET_OWNER", litellm_params, optional_params)
+
+
+def _resolve_s3_setting(
+    param_name: str,
+    env_var: str,
+    litellm_params: Mapping[str, object],
+    optional_params: Mapping[str, object] | None,
+) -> str | None:
     candidates: Final = tuple(
-        source.get("s3_encryption_key_id") for source in (litellm_params, optional_params) if source is not None
+        source.get(param_name) for source in (litellm_params, optional_params) if source is not None
     )
     explicit: Final = next((value for value in candidates if isinstance(value, str) and value), None)
-    return explicit or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID")
+    return explicit or get_secret_str(env_var)
 
 
 class CommonBatchFilesUtils:

litellm/types/llms/bedrock.py

@@ -4,7 +4,7 @@ from enum import Enum
 from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias
 
 from pydantic import BaseModel, ConfigDict
-from typing_extensions import ReadOnly, Required, TypedDict, override
+from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, override
 
 from .openai import ChatCompletionToolCallChunk
 
@@ -1082,6 +1082,7 @@ class BedrockS3InputDataConfig(TypedDict):
     """S3 input data configuration for Bedrock batch jobs."""
 
     s3Uri: str
+    s3BucketOwner: NotRequired[ReadOnly[str]]
 
 
 class BedrockInputDataConfig(TypedDict):
@@ -1095,6 +1096,7 @@ class BedrockS3OutputDataConfig(TypedDict, total=False):
 
     s3Uri: str
     s3EncryptionKeyId: str | None
+    s3BucketOwner: ReadOnly[str]
 
 
 class BedrockOutputDataConfig(TypedDict):

litellm/types/router.py

@@ -306,6 +306,7 @@ class CredentialLiteLLMParams(BaseModel):
     s3_endpoint_url: str | None = None
     s3_region_name: str | None = None
     s3_encryption_key_id: str | None = None
+    s3_bucket_owner: str | None = None
     aws_batch_role_arn: str | None = None
     s3_output_bucket_name: str | None = None
     bedrock_tags: list | None = None

litellm/types/utils.py

@@ -3830,6 +3830,7 @@ bedrock_batch_litellm_params: Final = (
     "s3_region_name",
     "s3_endpoint_url",
     "s3_output_bucket_name",
+    "s3_bucket_owner",
     "bedrock_tags",
 )
 

tests/test_litellm/llms/bedrock/batches/test_transformation.py

@@ -185,6 +185,91 @@ def test_create_request_omits_kms_key_when_absent(config):
     assert "s3EncryptionKeyId" not in s3out
 
 
+def _signed_batch_request(config, litellm_params: dict, optional_params: dict) -> dict:
+    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://in-bucket/in.jsonl"},
+            optional_params=optional_params,
+            litellm_params={"aws_batch_role_arn": "arn:aws:iam::1:role/r", **litellm_params},
+        )
+    return mock_sign.call_args.kwargs["data"]
+
+
+@pytest.mark.parametrize(
+    ("litellm_params", "optional_params", "env_owner", "expected_owner"),
+    [
+        pytest.param({"s3_bucket_owner": "111111111111"}, {}, None, "111111111111", id="litellm_params"),
+        pytest.param({}, {"s3_bucket_owner": "222222222222"}, None, "222222222222", id="optional_params"),
+        pytest.param({}, {}, "333333333333", "333333333333", id="env"),
+        pytest.param(
+            {"s3_bucket_owner": "111111111111"},
+            {"s3_bucket_owner": "222222222222"},
+            "333333333333",
+            "111111111111",
+            id="litellm_params_wins",
+        ),
+        pytest.param(
+            {}, {"s3_bucket_owner": "222222222222"}, "333333333333", "222222222222", id="optional_params_beats_env"
+        ),
+    ],
+)
+def test_create_request_sets_s3_bucket_owner_on_input_and_output(
+    config, monkeypatch, litellm_params, optional_params, env_owner, expected_owner
+):
+    monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False)
+    if env_owner is None:
+        monkeypatch.delenv("AWS_S3_BUCKET_OWNER", raising=False)
+    else:
+        monkeypatch.setenv("AWS_S3_BUCKET_OWNER", env_owner)
+
+    bedrock_request = _signed_batch_request(config, litellm_params, optional_params)
+
+    assert bedrock_request["inputDataConfig"] == {
+        "s3InputDataConfig": {"s3Uri": "s3://in-bucket/in.jsonl", "s3BucketOwner": expected_owner}
+    }
+    assert bedrock_request["outputDataConfig"] == {
+        "s3OutputDataConfig": {
+            "s3Uri": "s3://in-bucket/litellm-batch-outputs/litellm-batch-1/",
+            "s3BucketOwner": expected_owner,
+        }
+    }
+
+
+def test_create_request_omits_s3_bucket_owner_when_unset(config, monkeypatch):
+    monkeypatch.delenv("AWS_S3_BUCKET_OWNER", raising=False)
+    monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False)
+
+    bedrock_request = _signed_batch_request(config, {}, {})
+
+    assert bedrock_request["inputDataConfig"] == {"s3InputDataConfig": {"s3Uri": "s3://in-bucket/in.jsonl"}}
+    assert bedrock_request["outputDataConfig"] == {
+        "s3OutputDataConfig": {"s3Uri": "s3://in-bucket/litellm-batch-outputs/litellm-batch-1/"}
+    }
+
+
+def test_create_request_keeps_kms_key_alongside_s3_bucket_owner(config, monkeypatch):
+    monkeypatch.delenv("AWS_S3_BUCKET_OWNER", raising=False)
+    monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False)
+
+    bedrock_request = _signed_batch_request(
+        config, {"s3_bucket_owner": "111111111111", "s3_encryption_key_id": "kms-key-123"}, {}
+    )
+
+    assert bedrock_request["outputDataConfig"] == {
+        "s3OutputDataConfig": {
+            "s3Uri": "s3://in-bucket/litellm-batch-outputs/litellm-batch-1/",
+            "s3BucketOwner": "111111111111",
+            "s3EncryptionKeyId": "kms-key-123",
+        }
+    }
+
+
 def test_create_request_missing_input_file_id_raises(config):
     with pytest.raises(ValueError, match="input_file_id is required"):
         config.transform_create_batch_request(

tests/test_litellm/test_router.py

@@ -6294,6 +6294,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields():
                     "s3_bucket_name": "my-batch-bucket",
                     "s3_region_name": "us-east-1",
                     "s3_encryption_key_id": "arn:aws:kms:us-west-2:123:key/abc",
+                    "s3_bucket_owner": "111111111111",
                     "aws_batch_role_arn": "arn:aws:iam::123:role/batch-role",
                 },
             }
@@ -6311,6 +6312,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields():
     assert credentials["s3_bucket_name"] == "my-batch-bucket"
     assert credentials["s3_region_name"] == "us-east-1"
     assert credentials["s3_encryption_key_id"] == "arn:aws:kms:us-west-2:123:key/abc"
+    assert credentials["s3_bucket_owner"] == "111111111111"
     assert credentials["aws_batch_role_arn"] == "arn:aws:iam::123:role/batch-role"
 
 

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -31209,6 +31209,8 @@ export interface components {
             rpm?: number | null;
             /** S3 Bucket Name */
             s3_bucket_name?: string | null;
+            /** S3 Bucket Owner */
+            s3_bucket_owner?: string | null;
             /** S3 Encryption Key Id */
             s3_encryption_key_id?: string | null;
             /** S3 Endpoint Url */
@@ -42003,6 +42005,8 @@ export interface components {
             rpm?: number | null;
             /** S3 Bucket Name */
             s3_bucket_name?: string | null;
+            /** S3 Bucket Owner */
+            s3_bucket_owner?: string | null;
             /** S3 Encryption Key Id */
             s3_encryption_key_id?: string | null;
             /** S3 Endpoint Url */