[Feat] Bedrock Batches - Add support for custom KMS encryption keys in Bedrock Batch operations (#16662)

* add s3_encryption_key_id

* add s3EncryptionKeyId to BedrockS3OutputDataConfig

* use s3EncryptionKeyId in bedrock output

* docs s3_encryption_key_id

* test_bedrock_batch_with_encryption_key_in_post_request

docs/my-website/docs/providers/bedrock_batches.md

@@ -40,6 +40,8 @@ model_list:
       s3_access_key_id: os.environ/AWS_ACCESS_KEY_ID
       s3_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
       aws_batch_role_arn: arn:aws:iam::888602223428:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV
+      # Optional: Custom KMS encryption key for S3 output
+      # s3_encryption_key_id: arn:aws:kms:us-west-2:123456789012:key/12345678-1234-1234-1234-123456789012
     model_info: 
       mode: batch # 👈 SPECIFY MODE AS BATCH, to tell user this is a batch model
 ```
@@ -55,6 +57,12 @@ model_list:
 | `aws_batch_role_arn` | IAM role ARN for Bedrock batch operations. Bedrock Batch APIs require an IAM role ARN to be set. |
 | `mode: batch` | Indicates to LiteLLM this is a batch model |
 
+**Optional Parameters:**
+
+| Parameter | Description |
+|-----------|-------------|
+| `s3_encryption_key_id` | Custom KMS encryption key ID for S3 output data. If not specified, Bedrock uses AWS managed encryption keys. |
+
 ### 2. Create Virtual Key
 
 ```bash showLineNumbers title="create_virtual_key.sh"
@@ -174,6 +182,29 @@ When a `target_model_names` is specified, the file is written to the S3 bucket c
 
 LiteLLM only supports Bedrock Anthropic Models for Batch API. If you want other bedrock models file an issue [here](https://github.com/BerriAI/litellm/issues/new/choose).
 
+### How do I use a custom KMS encryption key?
+
+If your S3 bucket requires a custom KMS encryption key, you can specify it in your configuration using `s3_encryption_key_id`. This is useful for enterprise customers with specific encryption requirements.
+
+You can set the encryption key in 2 ways:
+
+1. **In config.yaml** (recommended):
+```yaml
+model_list:
+  - model_name: "bedrock-batch-claude"
+    litellm_params:
+      model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0
+      s3_encryption_key_id: arn:aws:kms:us-west-2:123456789012:key/12345678-1234-1234-1234-123456789012
+      # ... other params
+```
+
+2. **As an environment variable**:
+```bash
+export AWS_S3_ENCRYPTION_KEY_ID=arn:aws:kms:us-west-2:123456789012:key/12345678-1234-1234-1234-123456789012
+```
+
+
+
 ## Further Reading
 
 - [AWS Bedrock Batch Inference Documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/batch-inference.html)

litellm/llms/bedrock/batches/transformation.py

@@ -6,6 +6,7 @@ from httpx import Headers, Response
 
 from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.bedrock import (
     BedrockCreateBatchRequest,
     BedrockCreateBatchResponse,
@@ -140,10 +141,20 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
         }
         
         # Build output data config
+        s3_output_config: BedrockS3OutputDataConfig = BedrockS3OutputDataConfig(
+            s3Uri=f"s3://{output_bucket}/{output_key}"
+        )
+        
+        # Add optional KMS encryption key ID if provided
+        s3_encryption_key_id = (
+            litellm_params.get("s3_encryption_key_id")
+            or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID")
+        )
+        if s3_encryption_key_id:
+            s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id
+        
         output_data_config: BedrockOutputDataConfig = {
-            "s3OutputDataConfig": BedrockS3OutputDataConfig(
-                s3Uri=f"s3://{output_bucket}/{output_key}"
-            )
+            "s3OutputDataConfig": s3_output_config
         }
         
         # Create Bedrock batch request with proper typing

litellm/types/llms/bedrock.py

@@ -679,10 +679,11 @@ class BedrockInputDataConfig(TypedDict):
     s3InputDataConfig: BedrockS3InputDataConfig
 
 
-class BedrockS3OutputDataConfig(TypedDict):
+class BedrockS3OutputDataConfig(TypedDict, total=False):
     """S3 output data configuration for Bedrock batch jobs."""
 
     s3Uri: str
+    s3EncryptionKeyId: Optional[str]
 
 
 class BedrockOutputDataConfig(TypedDict):

litellm/types/router.py

@@ -205,6 +205,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
 
     # Batch/File API Params
     s3_bucket_name: Optional[str] = None
+    s3_encryption_key_id: Optional[str] = None
     gcs_bucket_name: Optional[str] = None
 
     # Vector Store Params
@@ -262,6 +263,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
         auto_router_embedding_model: Optional[str] = None,
         # Batch/File API Params
         s3_bucket_name: Optional[str] = None,
+        s3_encryption_key_id: Optional[str] = None,
         gcs_bucket_name: Optional[str] = None,
         **params,
     ):

tests/batches_tests/test_bedrock_files_and_batches.py

@@ -193,3 +193,53 @@ async def test_bedrock_retrieve_batch():
         assert batch_response.input_file_id == "s3://test-bucket/input/test-input.jsonl"
         assert batch_response.output_file_id == "s3://test-bucket/output/"
 
+
+def test_bedrock_batch_with_encryption_key_in_post_request():
+    """
+    Test that s3_encryption_key_id is included in the AWS POST request payload.
+    """
+    import json
+    import litellm
+    
+    test_kms_key_id = "arn:aws:kms:us-west-2:123456789012:key/12345678-1234-1234-1234-123456789012"
+    
+    captured_request_body = None
+    
+    def mock_post(*args, **kwargs):
+        nonlocal captured_request_body
+        if "data" in kwargs:
+            captured_request_body = kwargs["data"]
+        
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "jobArn": "arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/test-job",
+            "jobName": "test-job",
+            "status": "Submitted"
+        }
+        mock_response.status_code = 200
+        mock_response.raise_for_status.return_value = None
+        return mock_response
+    
+    with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post", side_effect=mock_post):
+        response = litellm.create_batch(
+            completion_window="24h",
+            endpoint="/v1/chat/completions",
+            input_file_id="s3://test-bucket/input/test.jsonl",
+            custom_llm_provider="bedrock",
+            model="us.anthropic.claude-3-5-sonnet-20240620-v1:0",
+            s3_encryption_key_id=test_kms_key_id,
+            aws_batch_role_arn="arn:aws:iam::123456789012:role/test-role"
+        )
+    
+    assert captured_request_body is not None, "Request body was not captured"
+    
+    request_data = json.loads(captured_request_body)
+    print("REQUEST DATA to bedrock batch creation", json.dumps(request_data, indent=4))
+    
+    assert "outputDataConfig" in request_data
+    assert "s3OutputDataConfig" in request_data["outputDataConfig"]
+    assert "s3EncryptionKeyId" in request_data["outputDataConfig"]["s3OutputDataConfig"]
+    assert request_data["outputDataConfig"]["s3OutputDataConfig"]["s3EncryptionKeyId"] == test_kms_key_id
+    
+    print("SUCCESS: s3_encryption_key_id properly included in AWS POST request")
+