From cb9965535e62afabb19811908ee8f137889ec006 Mon Sep 17 00:00:00 2001 From: Dawei Gu Date: Mon, 31 Aug 2026 11:03:37 -0700 Subject: [PATCH 1/2] feat(bedrock): allow configuring batch job timeout beyond the 24h default Bedrock's CreateModelInvocationJob API accepts timeoutDurationInHours anywhere from 24 to 168 hours, but LiteLLM's create_batch only ever sends 24. The OpenAI-shaped completion_window field is typed Literal["24h"] across every provider since that mirrors OpenAI's real batch API, so any other value passed through it is silently dropped in the Bedrock transform and Bedrock falls back to its own 24h default. Add a Bedrock-specific bedrock_batch_timeout_hours override, read from litellm_params or optional_params the same way bedrock_tags already is, validated against Bedrock's 24-168h range. completion_window and its shared type are untouched, so this only affects Bedrock batch requests that opt in, and every other provider keeps its existing behavior. --- .../llms/bedrock/batches/transformation.py | 35 ++++- .../bedrock/batches/test_transformation.py | 122 ++++++++++++++++++ 2 files changed, 154 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 7729cdfdb0d..8e7ee5b5032 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -60,6 +60,28 @@ def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]: ) from e +BEDROCK_BATCH_TIMEOUT_MIN_HOURS: Final = 24 +BEDROCK_BATCH_TIMEOUT_MAX_HOURS: Final = 168 + +_BEDROCK_BATCH_TIMEOUT_HOURS_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(int) + + +def _validate_bedrock_batch_timeout_hours(raw_value: object) -> int: + try: + hours = _BEDROCK_BATCH_TIMEOUT_HOURS_ADAPTER.validate_python(raw_value, strict=True) + except ValidationError as e: + raise ValueError( + f"Invalid 'bedrock_batch_timeout_hours' value {raw_value!r}. Expected an int between " + f"{BEDROCK_BATCH_TIMEOUT_MIN_HOURS} and {BEDROCK_BATCH_TIMEOUT_MAX_HOURS}." + ) from e + if not (BEDROCK_BATCH_TIMEOUT_MIN_HOURS <= hours <= BEDROCK_BATCH_TIMEOUT_MAX_HOURS): + raise ValueError( + f"'bedrock_batch_timeout_hours' must be between {BEDROCK_BATCH_TIMEOUT_MIN_HOURS} and " + f"{BEDROCK_BATCH_TIMEOUT_MAX_HOURS} (Bedrock's supported range), got {hours}." + ) + return hours + + class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): """ Config for Bedrock Batches - handles batch job creation and management for Bedrock @@ -234,12 +256,19 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): bedrock_request["tags"] = _validate_bedrock_tags(bedrock_tags) # Add optional parameters if provided + config_batch_timeout_hours: Final = litellm_params.get("bedrock_batch_timeout_hours") + batch_timeout_hours: Final = ( + config_batch_timeout_hours + if config_batch_timeout_hours is not None + else optional_params.get("bedrock_batch_timeout_hours") + ) completion_window: Final = create_batch_data.get("completion_window") - if completion_window: + if batch_timeout_hours is not None: + bedrock_request["timeoutDurationInHours"] = _validate_bedrock_batch_timeout_hours(batch_timeout_hours) + elif completion_window == "24h": # Map OpenAI completion window to Bedrock timeout # OpenAI uses "24h", Bedrock expects timeout in hours - if completion_window == "24h": - bedrock_request["timeoutDurationInHours"] = 24 + bedrock_request["timeoutDurationInHours"] = 24 # For Bedrock, we need to return a pre-signed request with AWS auth headers # Use common utility for AWS signing diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 7e5716a7495..2a5f79d2b18 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -255,6 +255,128 @@ def test_create_request_no_timeout_for_non_24h_window(config): assert "timeoutDurationInHours" not in mock_sign.call_args.kwargs["data"] +def test_create_request_bedrock_batch_timeout_hours_from_litellm_params(config): # test-quality-ok: asserts the built request payload, only observable via the mocked signer's call args + 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://b/in.jsonl"}, + optional_params={}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "bedrock_batch_timeout_hours": 72, + }, + ) + assert mock_sign.call_args.kwargs["data"]["timeoutDurationInHours"] == 72 + + +def test_create_request_bedrock_batch_timeout_hours_from_optional_params(config): # test-quality-ok: asserts the built request payload, only observable via the mocked signer's call args + 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://b/in.jsonl"}, + optional_params={"bedrock_batch_timeout_hours": 100}, + litellm_params={"aws_batch_role_arn": "arn:aws:iam::1:role/r"}, + ) + assert mock_sign.call_args.kwargs["data"]["timeoutDurationInHours"] == 100 + + +def test_create_request_bedrock_batch_timeout_hours_litellm_params_takes_precedence_over_optional_params( # test-quality-ok: asserts the built request payload, only observable via the mocked signer's call args + config, +): + 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://b/in.jsonl"}, + optional_params={"bedrock_batch_timeout_hours": 48}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "bedrock_batch_timeout_hours": 168, + }, + ) + assert mock_sign.call_args.kwargs["data"]["timeoutDurationInHours"] == 168 + + +def test_create_request_bedrock_batch_timeout_hours_overrides_completion_window(config): # test-quality-ok: asserts the built request payload, only observable via the mocked signer's call args + 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://b/in.jsonl", + "completion_window": "24h", + }, + optional_params={}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "bedrock_batch_timeout_hours": 96, + }, + ) + assert mock_sign.call_args.kwargs["data"]["timeoutDurationInHours"] == 96 + + +def test_create_request_bedrock_batch_timeout_hours_absent_falls_back_to_completion_window( # test-quality-ok: asserts the built request payload, only observable via the mocked signer's call args + config, +): + 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://b/in.jsonl", + "completion_window": "24h", + }, + optional_params={}, + litellm_params={"aws_batch_role_arn": "arn:aws:iam::1:role/r"}, + ) + assert mock_sign.call_args.kwargs["data"]["timeoutDurationInHours"] == 24 + + +@pytest.mark.parametrize( + "bad_value", + [23, 169, "48", 24.5, True, [24]], +) +def test_create_request_rejects_invalid_bedrock_batch_timeout_hours(config, bad_value): + 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"{}") + with pytest.raises(ValueError, match="bedrock_batch_timeout_hours"): + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "bedrock_batch_timeout_hours": bad_value, + }, + ) + + def test_create_request_forwards_bedrock_tags_from_litellm_params(config): tags = [ {"key": "application", "value": "genai-proxy"}, From c39c7c5c9a4c2f838bdcbba16fda57a684cdcae4 Mon Sep 17 00:00:00 2001 From: Dawei Gu Date: Mon, 31 Aug 2026 11:55:09 -0700 Subject: [PATCH 2/2] fix(bedrock): register bedrock_batch_timeout_hours as a batch-only param Address review feedback on the batch timeout override: - Register bedrock_batch_timeout_hours in bedrock_batch_litellm_params and _BANNED_REQUEST_BODY_PARAMS, matching bedrock_tags. Without this, a deployment configured with the override that also serves non-batch Bedrock requests would leak the field into extra_body and get it rejected by Bedrock, and a caller could set it directly in the request body instead of only through admin-configured litellm_params. - Annotate the validated value in _validate_bedrock_batch_timeout_hours with Final, matching the repo's convention. - Shorten a few new test names and wrap their signatures so the test-quality-ok suppressions they need stay within the 120-char limit. --- litellm/llms/bedrock/batches/transformation.py | 2 +- litellm/proxy/auth/auth_utils.py | 1 + litellm/types/utils.py | 1 + .../llms/bedrock/batches/test_transformation.py | 16 +++++++++++----- 4 files changed, 14 insertions(+), 6 deletions(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 8e7ee5b5032..64bb38ed380 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -68,7 +68,7 @@ _BEDROCK_BATCH_TIMEOUT_HOURS_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(int) def _validate_bedrock_batch_timeout_hours(raw_value: object) -> int: try: - hours = _BEDROCK_BATCH_TIMEOUT_HOURS_ADAPTER.validate_python(raw_value, strict=True) + hours: Final = _BEDROCK_BATCH_TIMEOUT_HOURS_ADAPTER.validate_python(raw_value, strict=True) except ValidationError as e: raise ValueError( f"Invalid 'bedrock_batch_timeout_hours' value {raw_value!r}. Expected an int between " diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 9b1a6ba5aa7..9bb5aeee968 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -297,6 +297,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( # reachable with the deployment's shared AWS credentials. "aws_bedrock_project_id", "bedrock_tags", + "bedrock_batch_timeout_hours", # Provider-specific endpoint overrides that flow into the outbound # request via ``optional_params``. Same threat as ``api_base``: # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4bf8289d725..3e86e513b75 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3525,6 +3525,7 @@ bedrock_batch_litellm_params: Final = ( "s3_region_name", "s3_output_bucket_name", "bedrock_tags", + "bedrock_batch_timeout_hours", ) all_litellm_params = ( diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 2a5f79d2b18..af5694c01bb 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -255,7 +255,9 @@ def test_create_request_no_timeout_for_non_24h_window(config): assert "timeoutDurationInHours" not in mock_sign.call_args.kwargs["data"] -def test_create_request_bedrock_batch_timeout_hours_from_litellm_params(config): # test-quality-ok: asserts the built request payload, only observable via the mocked signer's call args +def test_create_request_bedrock_batch_timeout_hours_from_litellm_params( # test-quality-ok: like bedrock_tags + config, +): with patch.object( config.common_utils, "generate_unique_job_name", @@ -274,7 +276,9 @@ def test_create_request_bedrock_batch_timeout_hours_from_litellm_params(config): assert mock_sign.call_args.kwargs["data"]["timeoutDurationInHours"] == 72 -def test_create_request_bedrock_batch_timeout_hours_from_optional_params(config): # test-quality-ok: asserts the built request payload, only observable via the mocked signer's call args +def test_create_request_bedrock_batch_timeout_hours_from_optional_params( # test-quality-ok: like bedrock_tags + config, +): with patch.object( config.common_utils, "generate_unique_job_name", @@ -290,7 +294,7 @@ def test_create_request_bedrock_batch_timeout_hours_from_optional_params(config) assert mock_sign.call_args.kwargs["data"]["timeoutDurationInHours"] == 100 -def test_create_request_bedrock_batch_timeout_hours_litellm_params_takes_precedence_over_optional_params( # test-quality-ok: asserts the built request payload, only observable via the mocked signer's call args +def test_create_request_bedrock_batch_timeout_hours_prefers_litellm_params( # test-quality-ok: like bedrock_tags config, ): with patch.object( @@ -311,7 +315,9 @@ def test_create_request_bedrock_batch_timeout_hours_litellm_params_takes_precede assert mock_sign.call_args.kwargs["data"]["timeoutDurationInHours"] == 168 -def test_create_request_bedrock_batch_timeout_hours_overrides_completion_window(config): # test-quality-ok: asserts the built request payload, only observable via the mocked signer's call args +def test_create_request_bedrock_batch_timeout_hours_overrides_completion_window( # test-quality-ok: like bedrock_tags + config, +): with patch.object( config.common_utils, "generate_unique_job_name", @@ -333,7 +339,7 @@ def test_create_request_bedrock_batch_timeout_hours_overrides_completion_window( assert mock_sign.call_args.kwargs["data"]["timeoutDurationInHours"] == 96 -def test_create_request_bedrock_batch_timeout_hours_absent_falls_back_to_completion_window( # test-quality-ok: asserts the built request payload, only observable via the mocked signer's call args +def test_create_request_bedrock_batch_timeout_hours_default_completion_window( # test-quality-ok: like bedrock_tags config, ): with patch.object(