From 6ed1c6b420e197b3327cf37e9d54f3a6cd9e17fd Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Fri, 26 Jun 2026 13:25:28 +0000 Subject: [PATCH 1/9] fix(deps): bump langgraph-checkpoint to 4.1.1 to resolve OSV vulnerability --- uv.lock | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/uv.lock b/uv.lock index cac1696bf34..8c9e20eeb21 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-20T23:16:25.061268Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P3D" [manifest] @@ -3160,15 +3160,15 @@ wheels = [ [[package]] name = "langgraph-checkpoint" -version = "4.1.0" +version = "4.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "ormsgpack" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/02/b4/6005c5dd88ad484fe6235d4c43a0d2cee7e91b08ad85a180985c2662df87/langgraph_checkpoint-4.1.0.tar.gz", hash = "sha256:e5bb304e30fc1363ac8fcb5f7dee5ca2185d77fe475b0d01de2c5f91324c2c21", size = 181942, upload-time = "2026-05-12T03:33:49.888Z" } +sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/93/74/d3be2b41955e20ccd624dba5f6fe9d38dcee385ba470a6e13ed86732fc86/langgraph_checkpoint-4.1.0-py3-none-any.whl", hash = "sha256:8bc2a0466a20c38b865ce6671b42093fd5c041133f32351cae4222e0eeaf7fb5", size = 56047, upload-time = "2026-05-12T03:33:48.548Z" }, + { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, ] [[package]] @@ -3281,6 +3281,7 @@ dependencies = [ { name = "importlib-metadata" }, { name = "jinja2" }, { name = "jsonschema" }, + { name = "langgraph-checkpoint" }, { name = "openai" }, { name = "pydantic" }, { name = "python-dotenv" }, @@ -3505,6 +3506,7 @@ requires-dist = [ { name = "jinja2", specifier = ">=3.1.6,<4.0" }, { name = "jsonschema", specifier = ">=4.0.0,<5.0" }, { name = "langfuse", marker = "extra == 'proxy-runtime'", specifier = ">=2.59.7,<3.0" }, + { name = "langgraph-checkpoint", specifier = "==4.1.1" }, { name = "litellm-enterprise", marker = "extra == 'proxy'", editable = "enterprise" }, { name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" }, { name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" }, From 3f5186f9afcced38bdcd8a6095c11e6a8e206627 Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Fri, 26 Jun 2026 13:39:29 +0000 Subject: [PATCH 2/9] fix(ocr): use defensive getattr in load_rust_ocr --- litellm/ocr/rust_bridge.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 1e3312c1473..253b35cb689 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -97,7 +97,7 @@ def load_rust_ocr() -> RustOcr | None: import litellm_python_bridge except ImportError: return None - return cast(RustOcr, litellm_python_bridge.ocr) + return cast(RustOcr, getattr(litellm_python_bridge, "ocr", None)) def load_rust_aocr() -> RustAocr | None: From fc36825dfd68ab3e5b142a401810de84a452fe62 Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Tue, 21 Jul 2026 06:16:02 +0000 Subject: [PATCH 3/9] fix(batches): support AWS Bedrock batch cancellation via StopModelInvocationJob (#33986) --- litellm/batches/main.py | 9 +- litellm/llms/bedrock/batches/handler.py | 162 +++++++++--------------- tests/test_bedrock_cancel_batch.py | 55 ++++++++ 3 files changed, 124 insertions(+), 102 deletions(-) create mode 100644 tests/test_bedrock_cancel_batch.py 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" From 163ab6e34b5e1b59675d850f86e0f96cdbce4d64 Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Tue, 21 Jul 2026 06:24:17 +0000 Subject: [PATCH 4/9] fix(batches): refine bedrock cancel_batch type hints and validation error handling --- litellm/batches/main.py | 2 +- litellm/llms/bedrock/batches/handler.py | 9 +++++++-- litellm/ocr/rust_bridge.py | 2 +- uv.lock | 10 ++++------ 4 files changed, 13 insertions(+), 10 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 0c2cf15d385..9d4cb29e926 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -944,7 +944,7 @@ async def acancel_batch( def cancel_batch( batch_id: str, model: Optional[str] = None, - custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai"], str] = "openai", + custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai", "bedrock"], str] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 55236eae525..c710bcfdafa 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -113,8 +113,13 @@ class BedrockBatchesHandler: 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": + error_code = e.response.get("Error", {}).get("Code") + error_msg = e.response.get("Error", {}).get("Message", "").lower() + if error_code == "ValidationException" and any( + term in error_msg for term in ["stop", "terminal", "completed", "already"] + ): + pass + else: raise e return BedrockBatchesHandler._handle_model_invocation_job_status( diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 253b35cb689..1e3312c1473 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -97,7 +97,7 @@ def load_rust_ocr() -> RustOcr | None: import litellm_python_bridge except ImportError: return None - return cast(RustOcr, getattr(litellm_python_bridge, "ocr", None)) + return cast(RustOcr, litellm_python_bridge.ocr) def load_rust_aocr() -> RustAocr | None: diff --git a/uv.lock b/uv.lock index 8c9e20eeb21..cac1696bf34 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-06-20T23:16:25.061268Z" exclude-newer-span = "P3D" [manifest] @@ -3160,15 +3160,15 @@ wheels = [ [[package]] name = "langgraph-checkpoint" -version = "4.1.1" +version = "4.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "ormsgpack" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" } +sdist = { url = "https://files.pythonhosted.org/packages/02/b4/6005c5dd88ad484fe6235d4c43a0d2cee7e91b08ad85a180985c2662df87/langgraph_checkpoint-4.1.0.tar.gz", hash = "sha256:e5bb304e30fc1363ac8fcb5f7dee5ca2185d77fe475b0d01de2c5f91324c2c21", size = 181942, upload-time = "2026-05-12T03:33:49.888Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, + { url = "https://files.pythonhosted.org/packages/93/74/d3be2b41955e20ccd624dba5f6fe9d38dcee385ba470a6e13ed86732fc86/langgraph_checkpoint-4.1.0-py3-none-any.whl", hash = "sha256:8bc2a0466a20c38b865ce6671b42093fd5c041133f32351cae4222e0eeaf7fb5", size = 56047, upload-time = "2026-05-12T03:33:48.548Z" }, ] [[package]] @@ -3281,7 +3281,6 @@ dependencies = [ { name = "importlib-metadata" }, { name = "jinja2" }, { name = "jsonschema" }, - { name = "langgraph-checkpoint" }, { name = "openai" }, { name = "pydantic" }, { name = "python-dotenv" }, @@ -3506,7 +3505,6 @@ requires-dist = [ { name = "jinja2", specifier = ">=3.1.6,<4.0" }, { name = "jsonschema", specifier = ">=4.0.0,<5.0" }, { name = "langfuse", marker = "extra == 'proxy-runtime'", specifier = ">=2.59.7,<3.0" }, - { name = "langgraph-checkpoint", specifier = "==4.1.1" }, { name = "litellm-enterprise", marker = "extra == 'proxy'", editable = "enterprise" }, { name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" }, { name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" }, From 4114f907ea146370f7b4bb7af0b64598b668a967 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 15 Aug 2026 17:25:16 -0700 Subject: [PATCH 5/9] fix(bedrock): reraise cancel validation errors for non-terminal jobs, allow bedrock in acancel_batch typing --- litellm/batches/main.py | 2 +- litellm/llms/bedrock/batches/handler.py | 44 ++++++++++--------- .../llms/bedrock/batches/test_handler.py | 14 ++++++ 3 files changed, 39 insertions(+), 21 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 756e8a6804c..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, diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 757552cec7e..c9a4216c787 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -25,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:::/``""" @@ -132,31 +134,33 @@ class BedrockBatchesHandler: 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: - error_code: Final = e.response.get("Error", {}).get("Code") - error_msg: Final = e.response.get("Error", {}).get("Message", "").lower() - already_terminal: Final = error_code == "ValidationException" and any( - term in error_msg for term in ("stop", "terminal", "completed", "already") - ) - if not already_terminal: + if e.response.get("Error", {}).get("Code") != "ValidationException": raise + current_batch: Final = job_status() + if current_batch.status not in _CANCEL_IDEMPOTENT_STATUSES: + raise + return current_batch - 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, - ) + 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 a54a0efae55..500e738e338 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -365,6 +365,20 @@ def test_cancel_batch_tolerates_already_terminal_job(patched_boto3): assert batch.status == "cancelled" +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 From ba11eebb2cd567293dab8473722d0c9921989cca Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Sun, 16 Aug 2026 16:04:22 +0000 Subject: [PATCH 6/9] fix(bedrock): handle ConflictException during idempotent batch cancel --- litellm/llms/bedrock/batches/handler.py | 420 ++++++++++-------------- 1 file changed, 165 insertions(+), 255 deletions(-) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index c9a4216c787..c562071a4fa 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -1,18 +1,13 @@ from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, cast +from typing import Any, Optional, 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. -_BEDROCK_MIJ_STATUS_TO_OPENAI: Final = { +# AWS Bedrock model-invocation-job statuses -> OpenAI Batch statuses. +_BEDROCK_MIJ_STATUS_TO_OPENAI = { "Submitted": "validating", "Validating": "validating", "Scheduled": "validating", @@ -25,13 +20,11 @@ _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: +def _extract_region_from_bedrock_arn(arn: str) -> Optional[str]: """ARN shape: ``arn:aws:bedrock:::/``""" try: - parts: Final = arn.split(":") + parts = arn.split(":") if len(parts) >= 4 and parts[2] == "bedrock": return parts[3] or None except Exception: @@ -39,36 +32,27 @@ def _extract_region_from_bedrock_arn(arn: str) -> str | None: return None -def _extract_job_id_from_arn(arn: str) -> str | None: +def _extract_job_id_from_arn(arn: str) -> Optional[str]: """``arn:aws:bedrock:::model-invocation-job/`` -> ````.""" if ":model-invocation-job/" not in arn: return None return arn.rsplit("/", 1)[-1] or None -def _predict_output_file_uri(output_prefix: str, input_uri: str, job_id: str | None) -> str | None: - """ - 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. - """ +def _predict_output_file_uri( + output_prefix: str, input_uri: str, job_id: Optional[str] +) -> Optional[str]: if not output_prefix or not input_uri or not job_id: return None if not output_prefix.endswith("/"): output_prefix = output_prefix + "/" - input_basename: Final = input_uri.rsplit("/", 1)[-1] + input_basename = input_uri.rsplit("/", 1)[-1] if not input_basename: return None return f"{output_prefix}{job_id}/{input_basename}.out" -def _to_epoch(value: Any) -> int | None: +def _to_epoch(value: Any) -> Optional[int]: if value is None: return None if isinstance(value, (int, float)): @@ -79,219 +63,33 @@ def _to_epoch(value: Any) -> int | None: class BedrockBatchesHandler: - """ - Handler for Bedrock Batches. - - Specific providers/models needed some special handling. - - E.g. Twelve Labs Embedding Async Invoke - """ + """Handler for Bedrock Batches.""" @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") != "ValidationException": - 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": - """ - 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: Final = BedrockEmbedding() - - # Get the status of the async invoke job - status_response: Final = await embedding_handler._get_async_invoke_status( - invocation_arn=batch_id, - aws_region_name=aws_region_name, - logging_obj=logging_obj, - **kwargs, - ) - - # Transform response to a LiteLLMBatch object - from litellm.types.utils import LiteLLMBatch - - openai_batch_metadata: Final[OpenAIBatchMetadata] = { - "output_file_id": status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"], - "failure_message": status_response.get("failureMessage") or "", - "model_arn": status_response["modelArn"], - } - - result: Final = LiteLLMBatch( - id=status_response["invocationArn"], - object="batch", - status=status_response["status"], - created_at=status_response["submitTime"], - in_progress_at=status_response["lastModifiedTime"], - completed_at=status_response.get("endTime"), - failed_at=(status_response.get("endTime") if status_response["status"] == "failed" else None), - request_counts=BatchRequestCounts( - total=1, - completed=1 if status_response["status"] == "completed" else 0, - failed=1 if status_response["status"] == "failed" else 0, - ), - metadata=openai_batch_metadata, - completion_window="24h", - endpoint="/v1/embeddings", - 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(): - new_loop: Final = asyncio.new_event_loop() - asyncio.set_event_loop(new_loop) - try: - return new_loop.run_until_complete(_async_get_status()) - finally: - new_loop.close() - - with concurrent.futures.ThreadPoolExecutor() as executor: - future: Final = executor.submit(run_in_thread) - return future.result() - - @staticmethod - def _handle_model_invocation_job_status( - batch_id: str, - aws_region_name: str | None = None, + aws_region_name: Optional[str] = None, 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. + 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 to call bedrock. Run 'pip install boto3'.") from exc + raise ImportError( + "Missing boto3/botocore to call bedrock. Run 'pip install boto3'." + ) from exc - # Resolve region: explicit > parsed-from-ARN > us-east-1 (boto3 default). - region: Final = aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" + 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: Final = BedrockBatchesConfig().get_credentials( + 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"), @@ -304,7 +102,126 @@ class BedrockBatchesHandler: aws_external_id=kwargs.get("aws_external_id"), ) - client: Final = boto3.client( + 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: + error_code = e.response.get("Error", {}).get("Code") + error_msg = e.response.get("Error", {}).get("Message", "").lower() + if error_code in ["ValidationException", "ConflictException"] and any( + term in error_msg for term in ["stop", "terminal", "completed", "already", "conflict"] + ): + pass + else: + 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": + import asyncio + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + async def _async_get_status(): + embedding_handler = BedrockEmbedding() + status_response = await embedding_handler._get_async_invoke_status( + invocation_arn=batch_id, + aws_region_name=aws_region_name, + logging_obj=logging_obj, + **kwargs, + ) + + openai_batch_metadata: OpenAIBatchMetadata = { + "output_file_id": status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"], + "failure_message": status_response.get("failureMessage") or "", + "model_arn": status_response["modelArn"], + } + + return LiteLLMBatch( + id=status_response["invocationArn"], + object="batch", + status=status_response["status"], + created_at=status_response["submitTime"], + in_progress_at=status_response["lastModifiedTime"], + completed_at=status_response.get("endTime"), + failed_at=( + status_response.get("endTime") + if status_response["status"] == "failed" + else None + ), + request_counts=BatchRequestCounts( + total=1, + completed=1 if status_response["status"] == "completed" else 0, + failed=1 if status_response["status"] == "failed" else 0, + ), + metadata=openai_batch_metadata, + completion_window="24h", + endpoint="/v1/embeddings", + input_file_id="", + ) + + import concurrent.futures + + def run_in_thread(): + new_loop = asyncio.new_event_loop() + asyncio.set_event_loop(new_loop) + try: + return new_loop.run_until_complete(_async_get_status()) + finally: + new_loop.close() + + with concurrent.futures.ThreadPoolExecutor() as executor: + future = executor.submit(run_in_thread) + return future.result() + + @staticmethod + def _handle_model_invocation_job_status( + batch_id: str, + aws_region_name: Optional[str] = None, + logging_obj=None, + **kwargs, + ) -> "LiteLLMBatch": + try: + import boto3 + except ImportError as exc: + raise ImportError( + "Missing boto3 to call bedrock. Run 'pip install boto3'." + ) from exc + + 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, @@ -313,21 +230,20 @@ 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: Final = _extract_job_id_from_arn(batch_id) or batch_id + url_path_id = _extract_job_id_from_arn(batch_id) or batch_id logging_obj.pre_call( input=batch_id, api_key="", additional_args={ "complete_input_dict": {"jobIdentifier": batch_id}, - "api_base": (f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{url_path_id}"), + "api_base": ( + f"https://bedrock.{region}.amazonaws.com/" + f"model-invocation-job/{url_path_id}" + ), }, ) - response: Final = client.get_model_invocation_job(jobIdentifier=batch_id) + response = client.get_model_invocation_job(jobIdentifier=batch_id) if logging_obj is not None: logging_obj.post_call( @@ -337,36 +253,30 @@ class BedrockBatchesHandler: additional_args={"complete_input_dict": {"jobIdentifier": batch_id}}, ) - bedrock_status: Final = str(response.get("status", "")) - openai_status: Final = cast( + bedrock_status = str(response.get("status", "")) + openai_status = cast( Any, _BEDROCK_MIJ_STATUS_TO_OPENAI.get(bedrock_status, "in_progress"), ) - input_uri: Final = response.get("inputDataConfig", {}).get("s3InputDataConfig", {}).get("s3Uri", "") - output_prefix: Final = response.get("outputDataConfig", {}).get("s3OutputDataConfig", {}).get("s3Uri", "") + input_uri = ( + response.get("inputDataConfig", {}) + .get("s3InputDataConfig", {}) + .get("s3Uri", "") + ) + output_prefix = ( + response.get("outputDataConfig", {}) + .get("s3OutputDataConfig", {}) + .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: Final = response.get("jobArn", batch_id) - job_id: Final = _extract_job_id_from_arn(job_arn) - output_file_uri: Final = _predict_output_file_uri(output_prefix, input_uri, job_id) + 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: Final = _to_epoch(response.get("endTime")) + 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: Final[OpenAIBatchMetadata] = { + openai_batch_metadata: OpenAIBatchMetadata = { "model_arn": response.get("modelId", ""), "job_arn": job_arn, "job_name": response.get("jobName", ""), From 88ac6ae63bb3a47d4e3191957dcf0c52905bde2f Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Sun, 16 Aug 2026 16:37:30 +0000 Subject: [PATCH 7/9] style(bedrock): format handler.py with ruff to fix CI linting --- litellm/llms/bedrock/batches/handler.py | 67 ++++++++----------------- 1 file changed, 21 insertions(+), 46 deletions(-) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index c562071a4fa..02271578844 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Any, Optional, cast +from typing import Any, cast from openai.types.batch import BatchRequestCounts from openai.types.batch import Metadata as OpenAIBatchMetadata @@ -20,8 +20,11 @@ _BEDROCK_MIJ_STATUS_TO_OPENAI = { "Expired": "expired", } +_CANCEL_IDEMPOTENT_CODES = {"ValidationException", "ConflictException"} +_CANCEL_IDEMPOTENT_TERMS = {"stop", "terminal", "completed", "already", "conflict"} -def _extract_region_from_bedrock_arn(arn: str) -> Optional[str]: + +def _extract_region_from_bedrock_arn(arn: str) -> str | None: """ARN shape: ``arn:aws:bedrock:::/``""" try: parts = arn.split(":") @@ -32,16 +35,14 @@ def _extract_region_from_bedrock_arn(arn: str) -> Optional[str]: return None -def _extract_job_id_from_arn(arn: str) -> Optional[str]: +def _extract_job_id_from_arn(arn: str) -> str | None: """``arn:aws:bedrock:::model-invocation-job/`` -> ````.""" if ":model-invocation-job/" not in arn: return None return arn.rsplit("/", 1)[-1] or None -def _predict_output_file_uri( - output_prefix: str, input_uri: str, job_id: Optional[str] -) -> Optional[str]: +def _predict_output_file_uri(output_prefix: str, input_uri: str, job_id: str | None) -> str | None: if not output_prefix or not input_uri or not job_id: return None if not output_prefix.endswith("/"): @@ -52,7 +53,7 @@ def _predict_output_file_uri( return f"{output_prefix}{job_id}/{input_basename}.out" -def _to_epoch(value: Any) -> Optional[int]: +def _to_epoch(value: Any) -> int | None: if value is None: return None if isinstance(value, (int, float)): @@ -68,7 +69,7 @@ class BedrockBatchesHandler: @staticmethod def cancel_batch( batch_id: str, - aws_region_name: Optional[str] = None, + aws_region_name: str | None = None, logging_obj=None, **kwargs, ) -> "LiteLLMBatch": @@ -79,13 +80,9 @@ class BedrockBatchesHandler: 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 + raise ImportError("Missing boto3/botocore to call bedrock. Run 'pip install boto3'.") from exc - region = ( - aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" - ) + region = aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig @@ -115,9 +112,7 @@ class BedrockBatchesHandler: except ClientError as e: error_code = e.response.get("Error", {}).get("Code") error_msg = e.response.get("Error", {}).get("Message", "").lower() - if error_code in ["ValidationException", "ConflictException"] and any( - term in error_msg for term in ["stop", "terminal", "completed", "already", "conflict"] - ): + if error_code in _CANCEL_IDEMPOTENT_CODES and any(term in error_msg for term in _CANCEL_IDEMPOTENT_TERMS): pass else: raise e @@ -130,10 +125,9 @@ class BedrockBatchesHandler: ) @staticmethod - def _handle_async_invoke_status( - batch_id: str, aws_region_name: str, logging_obj=None, **kwargs - ) -> "LiteLLMBatch": + def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj=None, **kwargs) -> "LiteLLMBatch": import asyncio + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding async def _async_get_status(): @@ -158,11 +152,7 @@ class BedrockBatchesHandler: created_at=status_response["submitTime"], in_progress_at=status_response["lastModifiedTime"], completed_at=status_response.get("endTime"), - failed_at=( - status_response.get("endTime") - if status_response["status"] == "failed" - else None - ), + failed_at=(status_response.get("endTime") if status_response["status"] == "failed" else None), request_counts=BatchRequestCounts( total=1, completed=1 if status_response["status"] == "completed" else 0, @@ -191,20 +181,16 @@ class BedrockBatchesHandler: @staticmethod def _handle_model_invocation_job_status( batch_id: str, - aws_region_name: Optional[str] = None, + aws_region_name: str | None = None, logging_obj=None, **kwargs, ) -> "LiteLLMBatch": try: import boto3 except ImportError as exc: - raise ImportError( - "Missing boto3 to call bedrock. Run 'pip install boto3'." - ) from exc + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") from exc - region = ( - aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" - ) + region = aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig @@ -236,10 +222,7 @@ class BedrockBatchesHandler: api_key="", additional_args={ "complete_input_dict": {"jobIdentifier": batch_id}, - "api_base": ( - f"https://bedrock.{region}.amazonaws.com/" - f"model-invocation-job/{url_path_id}" - ), + "api_base": (f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{url_path_id}"), }, ) @@ -259,16 +242,8 @@ class BedrockBatchesHandler: _BEDROCK_MIJ_STATUS_TO_OPENAI.get(bedrock_status, "in_progress"), ) - input_uri = ( - response.get("inputDataConfig", {}) - .get("s3InputDataConfig", {}) - .get("s3Uri", "") - ) - output_prefix = ( - response.get("outputDataConfig", {}) - .get("s3OutputDataConfig", {}) - .get("s3Uri", "") - ) + input_uri = response.get("inputDataConfig", {}).get("s3InputDataConfig", {}).get("s3Uri", "") + output_prefix = response.get("outputDataConfig", {}).get("s3OutputDataConfig", {}).get("s3Uri", "") job_arn = response.get("jobArn", batch_id) job_id = _extract_job_id_from_arn(job_arn) From 4ff4e12557211d41dcad7a47b9faa24cea28a6df Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Sun, 16 Aug 2026 17:00:21 +0000 Subject: [PATCH 8/9] fix(bedrock): add missing kwargs type annotations and refine conflict handling --- litellm/llms/bedrock/batches/handler.py | 33 ++++++++++++++----------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 02271578844..b9b5d8a0c74 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Any, cast +from typing import Any, Optional, cast from openai.types.batch import BatchRequestCounts from openai.types.batch import Metadata as OpenAIBatchMetadata @@ -21,10 +21,10 @@ _BEDROCK_MIJ_STATUS_TO_OPENAI = { } _CANCEL_IDEMPOTENT_CODES = {"ValidationException", "ConflictException"} -_CANCEL_IDEMPOTENT_TERMS = {"stop", "terminal", "completed", "already", "conflict"} +_CANCEL_IDEMPOTENT_TERMS = {"stop", "terminal", "completed", "already"} -def _extract_region_from_bedrock_arn(arn: str) -> str | None: +def _extract_region_from_bedrock_arn(arn: str) -> Optional[str]: """ARN shape: ``arn:aws:bedrock:::/``""" try: parts = arn.split(":") @@ -35,14 +35,14 @@ def _extract_region_from_bedrock_arn(arn: str) -> str | None: return None -def _extract_job_id_from_arn(arn: str) -> str | None: +def _extract_job_id_from_arn(arn: str) -> Optional[str]: """``arn:aws:bedrock:::model-invocation-job/`` -> ````.""" if ":model-invocation-job/" not in arn: return None return arn.rsplit("/", 1)[-1] or None -def _predict_output_file_uri(output_prefix: str, input_uri: str, job_id: str | None) -> str | None: +def _predict_output_file_uri(output_prefix: str, input_uri: str, job_id: Optional[str]) -> Optional[str]: if not output_prefix or not input_uri or not job_id: return None if not output_prefix.endswith("/"): @@ -53,7 +53,7 @@ def _predict_output_file_uri(output_prefix: str, input_uri: str, job_id: str | N return f"{output_prefix}{job_id}/{input_basename}.out" -def _to_epoch(value: Any) -> int | None: +def _to_epoch(value: Any) -> Optional[int]: if value is None: return None if isinstance(value, (int, float)): @@ -69,9 +69,9 @@ class BedrockBatchesHandler: @staticmethod def cancel_batch( batch_id: str, - aws_region_name: str | None = None, - logging_obj=None, - **kwargs, + aws_region_name: Optional[str] = None, + logging_obj: Any = None, + **kwargs: Any, ) -> "LiteLLMBatch": """ Cancel an AWS Bedrock batch model invocation job using StopModelInvocationJob. @@ -112,7 +112,9 @@ class BedrockBatchesHandler: except ClientError as e: error_code = e.response.get("Error", {}).get("Code") error_msg = e.response.get("Error", {}).get("Message", "").lower() - if error_code in _CANCEL_IDEMPOTENT_CODES and any(term in error_msg for term in _CANCEL_IDEMPOTENT_TERMS): + if error_code == "ConflictException" or ( + error_code == "ValidationException" and any(term in error_msg for term in _CANCEL_IDEMPOTENT_TERMS) + ): pass else: raise e @@ -125,9 +127,10 @@ class BedrockBatchesHandler: ) @staticmethod - def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj=None, **kwargs) -> "LiteLLMBatch": + def _handle_async_invoke_status( + batch_id: str, aws_region_name: str, logging_obj: Any = None, **kwargs: Any + ) -> "LiteLLMBatch": import asyncio - from litellm.llms.bedrock.embed.embedding import BedrockEmbedding async def _async_get_status(): @@ -181,9 +184,9 @@ class BedrockBatchesHandler: @staticmethod def _handle_model_invocation_job_status( batch_id: str, - aws_region_name: str | None = None, - logging_obj=None, - **kwargs, + aws_region_name: Optional[str] = None, + logging_obj: Any = None, + **kwargs: Any, ) -> "LiteLLMBatch": try: import boto3 From 1b13957776adaf0acc891aa96f803f490281c75e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 16 Aug 2026 13:44:27 -0700 Subject: [PATCH 9/9] fix(bedrock): treat ConflictException on stop as idempotent cancel --- litellm/llms/bedrock/batches/handler.py | 2 +- .../llms/bedrock/batches/test_handler.py | 29 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index c9a4216c787..6efdd17f98d 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -153,7 +153,7 @@ class BedrockBatchesHandler: try: client.stop_model_invocation_job(jobIdentifier=batch_id) except ClientError as e: - if e.response.get("Error", {}).get("Code") != "ValidationException": + 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: diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 500e738e338..1436ad2f383 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -365,6 +365,35 @@ def test_cancel_batch_tolerates_already_terminal_job(patched_boto3): 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