fix(bedrock): pass SSE-KMS key through to the batch input-file S3 upload - #35148
Conversation
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
|
|
Greptile SummaryAdds SSE-KMS support for Bedrock batch input-file uploads.
Confidence Score: 5/5The PR appears safe to merge No blocking failure remains
|
| Filename | Overview |
|---|---|
| litellm/llms/bedrock/common_utils.py | Adds a shared resolver that preserves deployment, optional-parameter, and environment precedence for the S3 encryption key |
| litellm/llms/bedrock/batches/transformation.py | Reuses the shared KMS key resolver when constructing Bedrock batch output configuration |
| litellm/llms/bedrock/files/transformation.py | Adds configured SSE-KMS headers before signing Bedrock batch input-file S3 uploads |
| tests/test_litellm/llms/bedrock/batches/test_transformation.py | Updates the secret resolver patch target after moving key resolution into the shared utility |
| tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py | Covers signed SSE-KMS headers, environment fallback, and omission when no key is configured |
Reviews (3): Last reviewed commit: "test(router): drop the duplicate s3_encr..." | Re-trigger Greptile
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
…itellm_wt_35148_merge
The staging merge tightened the LIT002 ceiling, so the three mutable dict literals this branch added now breach it. Build the S3 request headers as MappingProxyType and resolve the encryption key from a tuple of sources.
test_get_deployment_credentials_with_provider_bedrock_batch_fields already covers s3_encryption_key_id on the base branch, and the new test passes with every production file in this branch reverted, so it guards nothing.
TLDR
Problem this solves:
POST /v1/filesfails and no batch can be createdHow it solves it:
aws:kmsencryption headersUser Flow
A platform team runs Bedrock batch jobs through the gateway, and every S3 bucket they are allowed to use denies writes unless the object is encrypted with a specific KMS key.
Before
The upload that has to happen before any batch exists is rejected by S3, so the team never gets a file id and never reaches the batch endpoint at all.
bedrock/...deployment to the gateway config, pointing it at the locked-down bucket and setting the KMS key ARN on it, then starts the gatewayPOST https://litellm-domain/v1/fileswithpurpose=batch,model=bedrock-batch, and a JSONL body of chat-completion requests403with S3'sAccessDeniedtext saying the bucket policy requires server-side encryption with the bucket's KMS key403POST https://litellm-domain/v1/batchesis never reachable and the team cannot run batch inference against any of their bucketsAfter
The same upload succeeds, the object lands encrypted with the configured key, and the batch runs to completion.
bedrock/...deployment to the gateway config, pointing it at the locked-down bucket and setting the KMS key ARN on it, then starts the gatewayPOST https://litellm-domain/v1/fileswithpurpose=batch,model=bedrock-batch, and a JSONL body of chat-completion requests200with afile-id, the long scrambled kind the gateway hands back for model-routed uploadsaws s3api head-objecton the bucket shows the new object withServerSideEncryption: aws:kmsandSSEKMSKeyIdequal to the ARN configured on the deploymentPOST https://litellm-domain/v1/batcheswith that file id, a 24h completion window, and the/v1/chat/completionsendpoint, and gets back200with abatch_idRelevant issues
Fixes #35135
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
s3_encryption_key_idwas only ever used for the batch job'soutputDataConfig; the input-file upload that LiteLLM performs first went out with no encryption headers at all, which is what a KMS-enforced bucket rejects.Two proxies on one machine, same config file, same curl, only the code differs. BEFORE is
litellm_internal_stagingat f6587fa on port 4917, AFTER is this branch at a2806d4 on port 4918. Everything downstream is real: a customer-managed KMS key, a bucket whose policy deniesPutObjectunless the object carries that exact key, reals3.us-west-2.amazonaws.com, and a real Bedrock batch of 110 records onus.anthropic.claude-haiku-4-5-20251001-v1:0billing real $. The steps below are numbered to match the User Flow steps above, one for one.Step 1 of both flows is the deployment the admin sets up. The bucket policy is what makes this a real test rather than a header-shape check:
{ "Version": "2012-10-17", "Statement": [ { "Sid": "DenyUnencryptedUploads", "Effect": "Deny", "Principal": "*", "Action": "s3:PutObject", "Resource": "arn:aws:s3:::litellm-kms-qa-35148/*", "Condition": {"StringNotEquals": {"s3:x-amz-server-side-encryption": "aws:kms"}} }, { "Sid": "DenyWrongKmsKey", "Effect": "Deny", "Principal": "*", "Action": "s3:PutObject", "Resource": "arn:aws:s3:::litellm-kms-qa-35148/*", "Condition": {"StringNotEquals": {"s3:x-amz-server-side-encryption-aws-kms-key-id": "arn:aws:kms:us-west-2:439158074652:key/fa719dd3-a738-4b20-940a-a8df971041cd"}} } ] }Config, identical for both proxies:
Before
Steps 2 and 3, the 110-record upload, refused by S3 with the bucket's encryption requirement quoted back:
Step 4, a smaller file and two other content types, same answer every time:
Step 5, with no file id in hand there is nothing to start a batch from:
Step 6 has no command to run: dropping the bucket's encryption requirement is the one move the team's security policy rules out.
After
Steps 2 and 3, the same file, the same bucket, the same curl, now a managed file id:
Step 4, the object really landed encrypted with the configured key, not with a bucket default:
Step 5, the batch the user came for, created from that file id:
Step 6, the job walks all the way to completion. Polled every 30s, first appearance of each status:
And the results land in the bucket, 110 for 110, one line per request in the original file, real tokens billed:
The status walk above is read straight from Bedrock because the gateway's own
GET /v1/batches/{batch_id}, which is where a user would normally watch it, answersAccessDeniedException ... The provided resource ARN is from a different accountfor a job the proxy itself just created. That is #36155, pre-existing and reproducing onlitellm_internal_staging:The other bug this run surfaced is #36156: each record's
body.modelhas to be the rawbedrock/...id, because an alias makes the upload write amodelInputwith noanthropic_versionand every record then fails at invoke time behind a job Bedrock still reports asCompleted. Neither is touched by this diff and both reproduce onlitellm_internal_staging, so both are filed rather than fixed here.Type
🐛 Bug Fix
Changes
BedrockFilesConfig._sign_s3_requestnow takes the resolved KMS key and, when there is one, addsx-amz-server-side-encryption: aws:kmsplusx-amz-server-side-encryption-aws-kms-key-idto the headers it signs, so both headers land inside SigV4SignedHeadersrather than being appended afterwards (which S3 would answer withSignatureDoesNotMatch).The key itself is resolved by a new
resolve_s3_encryption_key_idinlitellm/llms/bedrock/common_utils.py, shared with the batch-creation path that previously inlined the same precedence chain:litellm_params, thenoptional_params, thenAWS_S3_ENCRYPTION_KEY_ID.That is the whole diff.
s3_encryption_key_idalready survivesCredentialLiteLLMParamsonlitellm_internal_stagingthanks to #24548, so nothing here toucheslitellm/types/router.py.Docs for the widened meaning of the key are in BerriAI/litellm-docs#810: it now encrypts the input file LiteLLM uploads as well as the batch output, so the uploading identity needs
kms:GenerateDataKeyon it.Caveats
optional_paramstier ofresolve_s3_encryption_key_idis unreachable in production:BaseLLMHTTPHandler.create_fileandcreate_batchboth passoptional_params={}to the transform, so "the batch path now readsoptional_paramstoo" cannot change any outcome, and deleting that argument leaves the whole bedrock suite greenlitellm_paramsthen the env var onlyaws_batch_role_arnandbedrock_tagsare read fromoptional_paramsthe same unreachable way a few lines above, so singling this one out makes the file inconsistent, and the tier becomes live the moment the caller stops passing an empty dictaws_access_key_idandaws_region_namenever reach the S3 PUT, because_sign_s3_requestreads credentials and region exclusively from that same emptyoptional_params; onlys3_region_nameis honoured fromlitellm_params, so an upload configured with a deployment-specific key pair is signed by the server environment's identity insteadoptional_paramshandoff inBaseLLMHTTPHandler.create_filein its own PR, tracked in [Bug]: Bedrock batch S3 upload ignores the deployment's own AWS credentials and region #36155QA runbook
Needs an S3 bucket whose policy denies PutObject unless
x-amz-server-side-encryptionisaws:kmswith a specific key ARN, plus credentials withkms:GenerateDataKeyon that key.bedrock/...deployment at that bucket withs3_encryption_key_idset to the key ARNpurpose=batch,model=<deployment>and a batch JSONL body; expect a 200 with a managed file id (pre-fix this is a 403AccessDeniednaming the bucket's encryption requirement)aws s3api head-object; expectServerSideEncryption: aws:kmsandSSEKMSKeyIdequal to the configured ARNFinal Attestation
Link to Devin session: https://app.devin.ai/sessions/050f0554f9bc486ab14918b6ab10aea1