diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 20d38bbb77f..a5df9b78601 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -826,7 +826,7 @@ def list_batches( async def acancel_batch( batch_id: str, model: str | None = None, - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -872,7 +872,7 @@ async def acancel_batch( def cancel_batch( batch_id: str, model: str | None = None, - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] | str = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] | str = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -993,9 +993,14 @@ def cancel_batch( timeout=timeout, max_retries=optional_params.max_retries, ) + elif custom_llm_provider == "bedrock": + response = BedrockBatchesHandler.cancel_batch( + batch_id=batch_id, + **kwargs, + ) else: raise litellm.exceptions.BadRequestError( - message=f"LiteLLM doesn't support {custom_llm_provider} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.", + message=f"LiteLLM doesn't support {custom_llm_provider} for 'cancel_batch'. Only 'openai', 'azure', 'vertex_ai', and 'bedrock' are supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 1752a727347..6efdd17f98d 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -1,11 +1,14 @@ from datetime import datetime -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, cast from openai.types.batch import BatchRequestCounts from openai.types.batch import Metadata as OpenAIBatchMetadata from litellm.types.utils import LiteLLMBatch +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + # AWS Bedrock model-invocation-job statuses → OpenAI Batch statuses. # Mirrors the mapping used by `BedrockBatchesConfig.transform_create_batch_response` # so create / retrieve return consistent statuses. @@ -22,6 +25,8 @@ _BEDROCK_MIJ_STATUS_TO_OPENAI: Final = { "Expired": "expired", } +_CANCEL_IDEMPOTENT_STATUSES: Final = frozenset({"cancelling", "cancelled", "completed", "failed", "expired"}) + def _extract_region_from_bedrock_arn(arn: str) -> str | None: """ARN shape: ``arn:aws:bedrock:::/``""" @@ -82,6 +87,81 @@ class BedrockBatchesHandler: E.g. Twelve Labs Embedding Async Invoke """ + @staticmethod + def cancel_batch( + batch_id: str, + aws_region_name: str | None = None, + logging_obj: "LiteLLMLoggingObj | None" = None, + aws_access_key_id: str | None = None, + aws_secret_access_key: str | None = None, + aws_session_token: str | None = None, + aws_session_name: str | None = None, + aws_profile_name: str | None = None, + aws_role_name: str | None = None, + aws_web_identity_token: str | None = None, + aws_sts_endpoint: str | None = None, + aws_external_id: str | None = None, + **kwargs: object, # kwargs-ok: litellm.cancel_batch forwards arbitrary user kwargs verbatim + ) -> "LiteLLMBatch": + try: + import boto3 + from botocore.exceptions import ClientError + except ImportError as exc: + raise ImportError("Missing boto3/botocore to call bedrock. Run 'pip install boto3'.") from exc + + region: Final = aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" + + from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + + creds: Final = BedrockBatchesConfig().get_credentials( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_region_name=region, + aws_session_name=aws_session_name, + aws_profile_name=aws_profile_name, + aws_role_name=aws_role_name, + aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, + ) + + client: Final = boto3.client( + "bedrock", + region_name=region, + aws_access_key_id=creds.access_key, + aws_secret_access_key=creds.secret_key, + aws_session_token=creds.token, + ) + + def job_status() -> "LiteLLMBatch": + return BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=batch_id, + aws_region_name=region, + logging_obj=logging_obj, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_session_name=aws_session_name, + aws_profile_name=aws_profile_name, + aws_role_name=aws_role_name, + aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, + ) + + try: + client.stop_model_invocation_job(jobIdentifier=batch_id) + except ClientError as e: + if e.response.get("Error", {}).get("Code") not in ("ValidationException", "ConflictException"): + raise + current_batch: Final = job_status() + if current_batch.status not in _CANCEL_IDEMPOTENT_STATUSES: + raise + return current_batch + + return job_status() + @staticmethod def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj=None, **kwargs) -> "LiteLLMBatch": """ diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 18780ccce0f..1436ad2f383 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -336,3 +336,100 @@ def test_logging_url_uses_bare_id_when_only_id_passed(patched_boto3): assert pre_kwargs["additional_args"]["api_base"] == ( f"https://bedrock.us-west-2.amazonaws.com/model-invocation-job/{JOB_ID}" ) + + +def test_cancel_batch_stops_job_and_returns_mapped_status(patched_boto3): + fake_client, boto_client_factory = patched_boto3 + fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopping") + + batch = BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN) + + fake_client.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN) + _, kwargs = boto_client_factory.call_args + assert kwargs["region_name"] == "us-west-2" + assert batch.status == "cancelling" + + +def test_cancel_batch_tolerates_already_terminal_job(patched_boto3): + from botocore.exceptions import ClientError + + fake_client, _ = patched_boto3 + fake_client.stop_model_invocation_job.side_effect = ClientError( + {"Error": {"Code": "ValidationException", "Message": "Job is already in a terminal state"}}, + "StopModelInvocationJob", + ) + fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped") + + batch = BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN) + + assert batch.status == "cancelled" + + +def test_cancel_batch_tolerates_conflict_on_already_stopped_job(patched_boto3): + from botocore.exceptions import ClientError + + fake_client, _ = patched_boto3 + fake_client.stop_model_invocation_job.side_effect = ClientError( + {"Error": {"Code": "ConflictException", "Message": "Job cannot be stopped in its current state"}}, + "StopModelInvocationJob", + ) + fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped") + + batch = BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN) + + assert batch.status == "cancelled" + + +def test_cancel_batch_reraises_conflict_when_job_not_terminal(patched_boto3): + from botocore.exceptions import ClientError + + fake_client, _ = patched_boto3 + fake_client.stop_model_invocation_job.side_effect = ClientError( + {"Error": {"Code": "ConflictException", "Message": "Operation conflicts with current job state"}}, + "StopModelInvocationJob", + ) + fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="InProgress") + + with pytest.raises(ClientError): + BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN) + + +def test_cancel_batch_reraises_validation_error_when_job_not_terminal(patched_boto3): + from botocore.exceptions import ClientError + + fake_client, _ = patched_boto3 + fake_client.stop_model_invocation_job.side_effect = ClientError( + {"Error": {"Code": "ValidationException", "Message": "Cannot stop job in current state"}}, + "StopModelInvocationJob", + ) + fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="InProgress") + + with pytest.raises(ClientError): + BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN) + + +def test_cancel_batch_reraises_other_client_errors(patched_boto3): + from botocore.exceptions import ClientError + + fake_client, _ = patched_boto3 + fake_client.stop_model_invocation_job.side_effect = ClientError( + {"Error": {"Code": "AccessDeniedException", "Message": "not authorized"}}, + "StopModelInvocationJob", + ) + + with pytest.raises(ClientError): + BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN) + + fake_client.get_model_invocation_job.assert_not_called() + + +def test_litellm_cancel_batch_dispatches_to_bedrock(patched_boto3): + import litellm + + fake_client, _ = patched_boto3 + fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped") + + batch = litellm.cancel_batch(batch_id=JOB_ARN, custom_llm_provider="bedrock") + + fake_client.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN) + assert batch.status == "cancelled"