diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 973388ca5bd..e4001566b8c 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/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] = { diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 7e24292a87e..f1066643874 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/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: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 10082cf2373..f9eaef5e891 100644 --- a/litellm/types/llms/bedrock.py +++ b/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): diff --git a/litellm/types/router.py b/litellm/types/router.py index a75b4654cab..6b43573d0e2 100644 --- a/litellm/types/router.py +++ b/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 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5a80644347e..45556acdfb5 100644 --- a/litellm/types/utils.py +++ b/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", ) diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 7e5716a7495..eb08c19cbdf 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/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( diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 8310d30d90e..d20fdff894a 100644 --- a/tests/test_litellm/test_router.py +++ b/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" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 00a8d979120..2f0cdfc47b1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/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 */