diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 84f40bec9ca..69f36e88b77 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/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 diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 9350c9c935c..13f1b7094fe 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/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 diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 1f90552e869..965aad79892 100644 --- a/litellm/types/litellm_params.py +++ b/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) diff --git a/litellm/types/router.py b/litellm/types/router.py index 4e3e4d92b3c..6e4fee320c6 100644 --- a/litellm/types/router.py +++ b/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 diff --git a/tests/unit/llms/bedrock/batches/test_transformation.py b/tests/unit/llms/bedrock/batches/test_transformation.py index 5360ac00018..148f13e7ec5 100644 --- a/tests/unit/llms/bedrock/batches/test_transformation.py +++ b/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"}, diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a9b56568afb..64c41fdc7ec 100644 --- a/tests/unit/test_utils.py +++ b/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)) diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 86a75ec65b9..79b79eede1a 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/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",