Merge b66056356b into 523b87ec71

litellm/llms/bedrock/batches/transformation.py

@@ -83,6 +83,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: 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 "
+            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
+
+
 def titan_embedding_usage_from_batch_output(model_output: Mapping[str, object]) -> Usage | None:
     """Titan embedding batch lines report usage as a top-level inputTextTokenCount, not a usage block."""
     if "embedding" not in model_output and "embeddingsByType" not in model_output:
@@ -268,12 +290,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

litellm/proxy/auth/auth_utils.py

@@ -434,6 +434,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
     "anthropic_workspace_id",
     "anthropic-workspace-id",
     "bedrock_tags",
+    "bedrock_batch_timeout_hours",
     # Deployment-pinned AgentCore session: every new caller-chosen value provisions another microVM.
     "agentcore_runtime_session_id",
     # Provider-specific endpoint overrides that flow into the outbound

litellm/types/litellm_params.py

@@ -117,6 +117,7 @@ class BedrockBatchConnection:
     s3_secret_access_key: str | None = None
     s3_encryption_key_id: str | None = None
     bedrock_tags: Sequence[Mapping[str, str]] | None = None
+    bedrock_batch_timeout_hours: int | None = None
 
 
 @dataclass(frozen=True, slots=True, kw_only=True)

litellm/types/router.py

@@ -453,6 +453,7 @@ class CredentialLiteLLMParams(LiteLLMBaseModel):
     aws_batch_role_arn: str | None = None
     s3_output_bucket_name: str | None = None
     bedrock_tags: list | None = None
+    bedrock_batch_timeout_hours: int | None = None
     ## IBM WATSONX ##
     watsonx_region_name: str | None = None
 

tests/unit/llms/bedrock/batches/test_transformation.py

@@ -466,6 +466,134 @@ 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(  # test-quality-ok: like bedrock_tags
+    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={},
+            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(  # test-quality-ok: like bedrock_tags
+    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": 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_prefers_litellm_params(  # test-quality-ok: like bedrock_tags
+    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(  # test-quality-ok: like bedrock_tags
+    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",
+                "bedrock_batch_timeout_hours": 96,
+            },
+        )
+    assert mock_sign.call_args.kwargs["data"]["timeoutDurationInHours"] == 96
+
+
+def test_create_request_bedrock_batch_timeout_hours_default_completion_window(  # test-quality-ok: like bedrock_tags
+    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"},

tests/unit/test_utils.py

@@ -6097,10 +6097,14 @@ def test_bedrock_batch_params_never_reach_the_provider():
     is extra="allow" and preserves them into litellm_params for the batch and files
     transformations that read them.
     """
-    configured = {
-        field: ([{"key": "team", "value": "configured-value"}] if field == "bedrock_tags" else "configured-value")
-        for field in bedrock_batch_litellm_params
-    }
+    def _configured_value(field: str) -> object:
+        if field == "bedrock_tags":
+            return [{"key": "team", "value": "configured-value"}]
+        if field == "bedrock_batch_timeout_hours":
+            return 48
+        return "configured-value"
+
+    configured = {field: _configured_value(field) for field in bedrock_batch_litellm_params}
     kwargs = {"a_real_provider_specific_param": 1, **configured}
 
     non_default = get_non_default_completion_params(dict(kwargs))

tests/unit/types/test_litellm_params.py

@@ -108,6 +108,7 @@ CONNECTION_NAMES: Final = (
     "s3_secret_access_key",
     "s3_encryption_key_id",
     "bedrock_tags",
+    "bedrock_batch_timeout_hours",
     "anthropic_federation_rule_id",
     "anthropic_organization_id",
     "anthropic_service_account_id",
@@ -295,6 +296,7 @@ BEDROCK_BATCH_NAMES: Final = (
     "s3_secret_access_key",
     "s3_encryption_key_id",
     "bedrock_tags",
+    "bedrock_batch_timeout_hours",
 )
 
 ARTIFACT_NAMES: Final = ("self", "use_client", "model_config", "rust")
@@ -539,7 +541,11 @@ TYPE_HINT_NAMESPACE: Final[Mapping[str, object]] = {
 
 LEAF_SAMPLES: Final[Mapping[type, Mapping[str, object]]] = {
     litellm_params.ProviderConnection: {"api_key": "k", "request_timeout": 1.5},
-    litellm_params.BedrockBatchConnection: {"aws_batch_role_arn": "arn", "bedrock_tags": ({"k": "v"},)},
+    litellm_params.BedrockBatchConnection: {
+        "aws_batch_role_arn": "arn",
+        "bedrock_tags": ({"k": "v"},),
+        "bedrock_batch_timeout_hours": 168,
+    },
     litellm_params.AnthropicFederationConnection: {
         "anthropic_federation_rule_id": "fdrl_1",
         "anthropic_issuer_ttl_seconds": 300,
@@ -703,6 +709,7 @@ NAMES_SHARED_WITH_TYPED_MODELS: Final[Mapping[str, tuple[str, ...]]] = MappingPr
             "azure_password",
             "azure_scope",
             "azure_username",
+            "bedrock_batch_timeout_hours",
             "bedrock_tags",
             "client_id",
             "client_secret",