diff --git a/docs/my-website/docs/providers/bedrock_batches.md b/docs/my-website/docs/providers/bedrock_batches.md index c262eef0e8..a1116f4107 100644 --- a/docs/my-website/docs/providers/bedrock_batches.md +++ b/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) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 2f3d00dddd..a9bc1b26c8 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/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 diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index bc752dd26a..330308e179 100644 --- a/litellm/types/llms/bedrock.py +++ b/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): diff --git a/litellm/types/router.py b/litellm/types/router.py index 3801d5bb78..2bf126211c 100644 --- a/litellm/types/router.py +++ b/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, ): diff --git a/tests/batches_tests/test_bedrock_files_and_batches.py b/tests/batches_tests/test_bedrock_files_and_batches.py index d082ed41ea..6ae373995d 100644 --- a/tests/batches_tests/test_bedrock_files_and_batches.py +++ b/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") +