diff --git a/litellm/batches/main.py b/litellm/batches/main.py index f124882b5a4..0c2cf15d385 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -1087,9 +1087,16 @@ def cancel_batch( timeout=timeout, max_retries=optional_params.max_retries, ) + elif custom_llm_provider == "bedrock": + from litellm.llms.bedrock.batches.handler import BedrockBatchesHandler + + response = BedrockBatchesHandler.cancel_batch( + batch_id=batch_id, + **kwargs, + ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.".format( + message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', 'vertex_ai', and 'bedrock' are supported.".format( custom_llm_provider ), model="n/a", diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index c071f331337..55236eae525 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -6,9 +6,7 @@ from openai.types.batch import Metadata as OpenAIBatchMetadata from litellm.types.utils import LiteLLMBatch -# 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. +# AWS Bedrock model-invocation-job statuses -> OpenAI Batch statuses. _BEDROCK_MIJ_STATUS_TO_OPENAI = { "Submitted": "validating", "Validating": "validating", @@ -44,17 +42,6 @@ def _extract_job_id_from_arn(arn: str) -> Optional[str]: def _predict_output_file_uri( output_prefix: str, input_uri: str, job_id: Optional[str] ) -> Optional[str]: - """ - Compute the deterministic per-job result file URI Bedrock writes to. - - Bedrock lays results out as:: - - //.out - - We compute it client-side so OpenAI-style ``client.files.content(output_file_id)`` - works without an extra S3 ``ListObjectsV2`` round-trip. Returns ``None`` if we - don't have enough info; callers should fall back to the bare prefix. - """ if not output_prefix or not input_uri or not job_id: return None if not output_prefix.endswith("/"): @@ -76,40 +63,76 @@ def _to_epoch(value: Any) -> Optional[int]: class BedrockBatchesHandler: - """ - Handler for Bedrock Batches. + """Handler for Bedrock Batches.""" - Specific providers/models needed some special handling. + @staticmethod + def cancel_batch( + batch_id: str, + aws_region_name: Optional[str] = None, + logging_obj=None, + **kwargs, + ) -> "LiteLLMBatch": + """ + Cancel an AWS Bedrock batch model invocation job using StopModelInvocationJob. + """ + 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 - E.g. Twelve Labs Embedding Async Invoke - """ + region = ( + aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" + ) + + from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + + creds = BedrockBatchesConfig().get_credentials( + aws_access_key_id=kwargs.get("aws_access_key_id"), + aws_secret_access_key=kwargs.get("aws_secret_access_key"), + aws_session_token=kwargs.get("aws_session_token"), + aws_region_name=region, + aws_session_name=kwargs.get("aws_session_name"), + aws_profile_name=kwargs.get("aws_profile_name"), + aws_role_name=kwargs.get("aws_role_name"), + aws_web_identity_token=kwargs.get("aws_web_identity_token"), + aws_sts_endpoint=kwargs.get("aws_sts_endpoint"), + aws_external_id=kwargs.get("aws_external_id"), + ) + + client = 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, + ) + + try: + client.stop_model_invocation_job(jobIdentifier=batch_id) + except ClientError as e: + # Idempotency: if job is already Stopping/Stopped/Completed, swallow ValidationException + if e.response.get("Error", {}).get("Code") != "ValidationException": + raise e + + return BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=batch_id, + aws_region_name=region, + logging_obj=logging_obj, + **kwargs, + ) @staticmethod def _handle_async_invoke_status( batch_id: str, aws_region_name: str, logging_obj=None, **kwargs ) -> "LiteLLMBatch": - """ - Handle async invoke status check for AWS Bedrock. - - This is for Twelve Labs Embedding Async Invoke. - - Args: - batch_id: The async invoke ARN - aws_region_name: AWS region name - **kwargs: Additional parameters - - Returns: - dict: Status information including status, output_file_id (S3 URL), etc. - """ import asyncio - from litellm.llms.bedrock.embed.embedding import BedrockEmbedding async def _async_get_status(): - # Create embedding handler instance embedding_handler = BedrockEmbedding() - - # Get the status of the async invoke job status_response = await embedding_handler._get_async_invoke_status( invocation_arn=batch_id, aws_region_name=aws_region_name, @@ -117,18 +140,13 @@ class BedrockBatchesHandler: **kwargs, ) - # Transform response to a LiteLLMBatch object - from litellm.types.utils import LiteLLMBatch - openai_batch_metadata: OpenAIBatchMetadata = { - "output_file_id": status_response["outputDataConfig"][ - "s3OutputDataConfig" - ]["s3Uri"], + "output_file_id": status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"], "failure_message": status_response.get("failureMessage") or "", "model_arn": status_response["modelArn"], } - result = LiteLLMBatch( + return LiteLLMBatch( id=status_response["invocationArn"], object="batch", status=status_response["status"], @@ -151,10 +169,6 @@ class BedrockBatchesHandler: input_file_id="", ) - return result - - # Since this function is called from within an async context via run_in_executor, - # we need to create a new event loop in a thread to avoid conflicts import concurrent.futures def run_in_thread(): @@ -176,37 +190,6 @@ class BedrockBatchesHandler: logging_obj=None, **kwargs, ) -> "LiteLLMBatch": - """ - Handle ``GetModelInvocationJob`` status check for AWS Bedrock bulk batch - inference jobs (the ARN type returned by ``CreateModelInvocationJob``). - - ``CreateModelInvocationJob`` lives on the Bedrock **control plane** - (``bedrock..amazonaws.com``), distinct from the data-plane - ``bedrock-runtime`` endpoint that serves Twelve Labs async-invoke ARNs. - The two ARN families therefore can't share a handler — see - ``litellm/batches/main.py`` for the dispatch. - - Args: - batch_id: A ``arn:aws:bedrock:::model-invocation-job/`` - ARN (or just the trailing job id; both are accepted by - ``GetModelInvocationJob``). - aws_region_name: Region for the boto3 ``bedrock`` client. If omitted, - we fall back to parsing the region out of ``batch_id`` itself. - logging_obj: Optional litellm logging object. - **kwargs: Optional AWS credential overrides - (``aws_access_key_id``, ``aws_secret_access_key``, - ``aws_session_token``, ``aws_profile_name``, - ``aws_role_name``, ``aws_session_name``, - ``aws_web_identity_token``, ``aws_sts_endpoint``, - ``aws_external_id``). Unknown keys are ignored. - - Returns: - ``LiteLLMBatch`` shaped like an OpenAI Batch resource. Note that - ``request_counts`` is always ``(0, 0, 0)`` because - ``GetModelInvocationJob`` does not surface per-record counts; - callers that need accurate counts should parse - ``manifest.json.out`` from the output S3 prefix. - """ try: import boto3 except ImportError as exc: @@ -214,15 +197,10 @@ class BedrockBatchesHandler: "Missing boto3 to call bedrock. Run 'pip install boto3'." ) from exc - # Resolve region: explicit > parsed-from-ARN > us-east-1 (boto3 default). region = ( aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" ) - # Resolve credentials through the same path the rest of the bedrock - # provider uses, so model_list / env / role-assumption configs are - # honored. We instantiate BedrockBatchesConfig (which extends - # BaseAWSLLM) lazily to avoid a circular import at module load. from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig creds = BedrockBatchesConfig().get_credentials( @@ -247,10 +225,6 @@ class BedrockBatchesHandler: ) if logging_obj is not None: - # Use the bare job id in the logged URL so we don't double up the - # `model-invocation-job/` segment when `batch_id` is a full ARN. - # `GetModelInvocationJob` accepts either form, but only the bare id - # produces a sensible-looking URL in logs. url_path_id = _extract_job_id_from_arn(batch_id) or batch_id logging_obj.pre_call( input=batch_id, @@ -291,26 +265,12 @@ class BedrockBatchesHandler: .get("s3Uri", "") ) - # Bedrock returns the output *prefix* the user supplied at job creation. - # Actual results land at //.out — we - # surface that single-file URI as `output_file_id` so the OpenAI-style - # download flow works without an extra S3 listing call. We deliberately - # do NOT fall back to the bare prefix when prediction fails: a prefix - # is not a downloadable object, so handing it back as `output_file_id` - # would reproduce the very NoSuchKey bug this handler exists to fix. - # The bare prefix is preserved in metadata for callers that want the - # `manifest.json.out` or want to do their own listing. job_arn = response.get("jobArn", batch_id) job_id = _extract_job_id_from_arn(job_arn) output_file_uri = _predict_output_file_uri(output_prefix, input_uri, job_id) completed_at = _to_epoch(response.get("endTime")) - # Note: metadata uses "" (not None) for unknown URIs to satisfy the - # OpenAI Batch metadata schema, which is `dict[str, str]`. The - # `output_file_id` field on the LiteLLMBatch itself does carry None - # correctly (see below), so callers should branch on that, not on - # `metadata["output_file_uri"]`. openai_batch_metadata: OpenAIBatchMetadata = { "model_arn": response.get("modelId", ""), "job_arn": job_arn, diff --git a/tests/test_bedrock_cancel_batch.py b/tests/test_bedrock_cancel_batch.py new file mode 100644 index 00000000000..e52cc2c447e --- /dev/null +++ b/tests/test_bedrock_cancel_batch.py @@ -0,0 +1,55 @@ +from unittest.mock import MagicMock, patch +import pytest + +import litellm +from litellm.llms.bedrock.batches.handler import BedrockBatchesHandler + + +@patch("boto3.client") +def test_bedrock_cancel_batch_handler(mock_boto_client): + mock_client_instance = MagicMock() + mock_boto_client.return_value = mock_client_instance + + mock_client_instance.stop_model_invocation_job.return_value = {} + mock_client_instance.get_model_invocation_job.return_value = { + "jobArn": "arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + "status": "Stopping", + "submitTime": 1700000000, + "lastModifiedTime": 1700000100, + } + + res = BedrockBatchesHandler.cancel_batch( + batch_id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + aws_region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + + mock_client_instance.stop_model_invocation_job.assert_called_once_with( + jobIdentifier="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id" + ) + assert res.status == "cancelling" + + +@patch("boto3.client") +def test_litellm_cancel_batch_bedrock_dispatcher(mock_boto_client): + mock_client_instance = MagicMock() + mock_boto_client.return_value = mock_client_instance + + mock_client_instance.stop_model_invocation_job.return_value = {} + mock_client_instance.get_model_invocation_job.return_value = { + "jobArn": "arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + "status": "Stopped", + "submitTime": 1700000000, + "lastModifiedTime": 1700000100, + } + + res = litellm.cancel_batch( + batch_id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + custom_llm_provider="bedrock", + aws_region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + + assert res.status == "cancelled"