diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index f73d5f5d3cb..b408d2f620c 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast @@ -6,6 +6,7 @@ from openai.types.batch import BatchRequestCounts from openai.types.batch import Metadata as OpenAIBatchMetadata from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.types.llms.bedrock import AwsSessionTag from litellm.types.utils import LiteLLMBatch if TYPE_CHECKING: @@ -116,6 +117,7 @@ class BedrockBatchesHandler: aws_web_identity_token: str | None = None, aws_sts_endpoint: str | None = None, aws_external_id: str | None = None, + aws_session_tags: Sequence[AwsSessionTag] | None = None, **kwargs: object, # kwargs-ok: litellm.cancel_batch forwards arbitrary user kwargs verbatim ) -> "LiteLLMBatch": try: @@ -139,6 +141,7 @@ class BedrockBatchesHandler: aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) client: Final = boto3.client( @@ -163,6 +166,7 @@ class BedrockBatchesHandler: aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) try: diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 4ccddc33e77..03daafcad72 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -482,10 +482,36 @@ def test_litellm_cancel_batch_dispatches_to_bedrock(patched_boto3): assert batch.status == "cancelled" +class _TagGatedSTSClient: + """Stands in for STS behind a trust policy that only admits sessions carrying ``tags``.""" + + def __init__(self, tags: list[dict[str, str]], access_key_id: str) -> None: + self._tags = tags + self._access_key_id = access_key_id + + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + from botocore.exceptions import ClientError + + if list(params.get("Tags", ())) != self._tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": self._access_key_id, + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.now(timezone.utc) + timedelta(minutes=30), + } + } + + def test_handle_model_invocation_job_status_builds_the_client_from_the_tagged_session(monkeypatch): """Status polling must assume the role with the deployment's session tags, like every other call.""" - from botocore.exceptions import ClientError - monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) monkeypatch.delenv("AWS_ROLE_ARN", raising=False) tags = [{"Key": "team", "Value": "genai"}] @@ -493,28 +519,9 @@ def test_handle_model_invocation_job_status_builds_the_client_from_the_tagged_se fake_bedrock = MagicMock() fake_bedrock.get_model_invocation_job.return_value = _fake_boto3_response() - class FakeSTSClient: - def get_caller_identity(self): - return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} - - def assume_role(self, **params): - if list(params.get("Tags", ())) != tags: - raise ClientError( - {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, - "AssumeRole", - ) - return { - "Credentials": { - "AccessKeyId": "ASIABATCHSTATUSTAGGED", - "SecretAccessKey": "assumed-secret", - "SessionToken": "assumed-session-token", - "Expiration": datetime.now(timezone.utc) + timedelta(minutes=30), - } - } - def boto3_client(service_name, **kwargs): if service_name == "sts": - return FakeSTSClient() + return _TagGatedSTSClient(tags, "ASIABATCHSTATUSTAGGED") bedrock_client_kwargs.append(kwargs) return fake_bedrock @@ -530,3 +537,36 @@ def test_handle_model_invocation_job_status_builds_the_client_from_the_tagged_se assert batch.status == "completed" assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHSTATUSTAGGED"] + + +def test_cancel_batch_stops_and_polls_the_job_with_the_tagged_session(monkeypatch): + """Cancelling on a tag-gated role must forward the deployment's session tags to both the stop and status calls.""" + import litellm + + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + bedrock_client_kwargs: list[dict] = [] + fake_bedrock = MagicMock() + fake_bedrock.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped") + + def boto3_client(service_name, **kwargs): + if service_name == "sts": + return _TagGatedSTSClient(tags, "ASIABATCHCANCELTAGGED") + bedrock_client_kwargs.append(kwargs) + return fake_bedrock + + with patch("boto3.client", side_effect=boto3_client): + batch = litellm.cancel_batch( + batch_id=JOB_ARN, + custom_llm_provider="bedrock", + aws_access_key_id="AKIABATCHCANCELCALLER", + aws_secret_access_key="pod-caller-secret", + aws_role_name="arn:aws:iam::999999999999:role/litellm-batch-role", + aws_session_name="litellm-batch-session", + aws_session_tags=tags, + ) + + fake_bedrock.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN) + assert batch.status == "cancelled" + assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHCANCELTAGGED"] * 2