From d4ceb90cfa805acfe082ef31054486a7e371895a Mon Sep 17 00:00:00 2001 From: Dawei Gu Date: Thu, 30 Apr 2026 11:32:00 -0700 Subject: [PATCH] fix(bedrock): address PR feedback on model-invocation-job retrieve Addresses Greptile P2 findings on #26834: 1. Use the bare job id (not the full ARN) when constructing the `api_base` URL for `pre_call` logging. Passing the full ARN double- counts the `model-invocation-job/` segment and embeds colons in the path, producing misleading log lines. 2. Drop the `or output_prefix` fallback when `_predict_output_file_uri` returns None. A bare prefix is not a downloadable object and surfacing it as `output_file_id` re-creates the very NoSuchKey bug this handler exists to fix. The bare prefix is still preserved in `metadata["output_s3_uri"]` for callers that want to do their own S3 listing or read `manifest.json.out`. `metadata["output_file_uri"]` uses "" rather than None to satisfy the OpenAI Batch metadata schema (`dict[str, str]`); callers should branch on the typed `output_file_id` field instead. Also expands test coverage on the new code path: - new "stay None" regression test for the prediction-fail case - pre_call/post_call logging hook assertions (incl. the bare-id URL) - explicit cancelled_at / expired_at coverage - _to_epoch type-handling matrix and the boto3 ImportError branch - defensive _extract_region_from_bedrock_arn exception path - empty-basename case for _predict_output_file_uri Patch coverage on the changed lines is now 100% (the only remaining uncovered lines in the file belong to the pre-existing `_handle_async_invoke_status` method, which this PR does not touch). Made-with: Cursor --- litellm/llms/bedrock/batches/handler.py | 26 ++- .../llms/bedrock/batches/test_handler.py | 157 ++++++++++++++++++ 2 files changed, 176 insertions(+), 7 deletions(-) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index cb9dcfab893..c071f331337 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -247,6 +247,11 @@ 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, api_key="", @@ -254,7 +259,7 @@ class BedrockBatchesHandler: "complete_input_dict": {"jobIdentifier": batch_id}, "api_base": ( f"https://bedrock.{region}.amazonaws.com/" - f"model-invocation-job/{batch_id}" + f"model-invocation-job/{url_path_id}" ), }, ) @@ -289,16 +294,23 @@ class BedrockBatchesHandler: # 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. The bare - # prefix is preserved in metadata for callers that want the manifest. + # 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) or output_prefix - ) + 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, @@ -306,7 +318,7 @@ class BedrockBatchesHandler: "failure_message": response.get("message") or "", "input_s3_uri": input_uri, "output_s3_uri": output_prefix, - "output_file_uri": output_file_uri, + "output_file_uri": output_file_uri or "", } return LiteLLMBatch( diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 2f88f96e9cd..18780ccce0f 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -22,6 +22,7 @@ from litellm.llms.bedrock.batches.handler import ( # noqa: E402 _extract_job_id_from_arn, _extract_region_from_bedrock_arn, _predict_output_file_uri, + _to_epoch, ) JOB_ID = "abc1234567" @@ -67,6 +68,42 @@ def test_extract_region_from_arn(): assert _extract_region_from_bedrock_arn("not-an-arn") is None +def test_extract_region_swallows_unexpected_split_errors(): + """Defensive `except Exception` branch — anything that isn't a plain str + should fall through to ``None`` rather than blow up.""" + + class WeirdArn: + def split(self, _sep): + raise RuntimeError("boom") + + assert _extract_region_from_bedrock_arn(WeirdArn()) is None # type: ignore[arg-type] + + +def test_predict_output_file_uri_returns_none_for_directory_input_uri(): + """Input URI ending in `/` has an empty basename — we must bail rather + than emit ``//.out``.""" + assert ( + _predict_output_file_uri(OUTPUT_PREFIX, "s3://bucket/inputs/", JOB_ID) is None + ) + + +_DT = datetime(2026, 4, 28, 12, 0, 0, tzinfo=timezone.utc) + + +@pytest.mark.parametrize( + "value,expected", + [ + (None, None), + (1730000000, 1730000000), + (1730000000.5, 1730000000), + (_DT, int(_DT.timestamp())), + ("2026-04-28T12:00:00Z", None), # strings aren't supported -> None + ], +) +def test_to_epoch_handles_supported_types(value, expected): + assert _to_epoch(value) == expected + + def test_extract_job_id_from_arn(): assert _extract_job_id_from_arn(JOB_ARN) == JOB_ID assert ( @@ -179,3 +216,123 @@ def test_failure_message_propagates(patched_boto3): assert batch.status == "failed" assert batch.failed_at == int(END_TIME.timestamp()) assert batch.metadata["failure_message"] == "Input file failed validation" + + +def test_completed_with_unpredictable_output_uri_stays_none(patched_boto3): + """ + Regression guard for the original NoSuchKey bug: if Bedrock's response is + missing pieces we need to compute the per-job output file path (here, the + input s3Uri), `output_file_id` must stay `None` rather than fall back to + the bare prefix. Falling back to the prefix is what produced the original + NoSuchKey error this PR fixes. + """ + fake_client, _ = patched_boto3 + incomplete_response = _fake_boto3_response(status="Completed") + incomplete_response["inputDataConfig"] = {"s3InputDataConfig": {"s3Uri": ""}} + fake_client.get_model_invocation_job.return_value = incomplete_response + + batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN) + + assert batch.status == "completed" + # output_file_id MUST be None (not the bare prefix) — that's the whole + # point of this regression test. Callers branch on this field. + assert batch.output_file_id is None + # The metadata field uses "" because OpenAI Batch metadata is dict[str, str]; + # callers should branch on `output_file_id` (above) instead. + assert batch.metadata["output_file_uri"] == "" + # The bare prefix is still preserved in metadata so callers can list it. + assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX + + +def test_cancelled_status_sets_cancelled_at(patched_boto3): + fake_client, _ = patched_boto3 + fake_client.get_model_invocation_job.return_value = _fake_boto3_response( + status="Stopped" + ) + + batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN) + + assert batch.status == "cancelled" + assert batch.cancelled_at == int(END_TIME.timestamp()) + assert batch.completed_at is None + assert batch.failed_at is None + assert batch.expired_at is None + + +def test_expired_status_sets_expired_at(patched_boto3): + fake_client, _ = patched_boto3 + fake_client.get_model_invocation_job.return_value = _fake_boto3_response( + status="Expired" + ) + + batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN) + + assert batch.status == "expired" + assert batch.expired_at == int(END_TIME.timestamp()) + assert batch.completed_at is None + assert batch.failed_at is None + assert batch.cancelled_at is None + + +def test_logging_obj_pre_and_post_call_invoked(patched_boto3): + """`pre_call` / `post_call` get called with sensible payloads when a + `logging_obj` is supplied.""" + _, _ = patched_boto3 + logging_obj = MagicMock() + + BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=JOB_ARN, logging_obj=logging_obj + ) + + logging_obj.pre_call.assert_called_once() + logging_obj.post_call.assert_called_once() + + pre_kwargs = logging_obj.pre_call.call_args.kwargs + assert pre_kwargs["input"] == JOB_ARN + assert pre_kwargs["additional_args"]["complete_input_dict"] == { + "jobIdentifier": JOB_ARN + } + # Logged URL must use the bare job id, not the full ARN, so it doesn't + # double the `model-invocation-job/` segment or embed colons in the path. + assert pre_kwargs["additional_args"]["api_base"] == ( + f"https://bedrock.us-west-2.amazonaws.com/model-invocation-job/{JOB_ID}" + ) + + post_kwargs = logging_obj.post_call.call_args.kwargs + assert post_kwargs["input"] == JOB_ARN + assert post_kwargs["original_response"]["jobArn"] == JOB_ARN + + +def test_missing_boto3_raises_helpful_import_error(): + """If boto3 isn't installed we should raise a clear, actionable + ImportError rather than letting a NameError escape.""" + real_import = ( + __builtins__["__import__"] + if isinstance(__builtins__, dict) + else __builtins__.__import__ + ) + + def fake_import(name, *args, **kwargs): + if name == "boto3": + raise ImportError("No module named 'boto3'") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=fake_import): + with pytest.raises(ImportError, match="pip install boto3"): + BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN) + + +def test_logging_url_uses_bare_id_when_only_id_passed(patched_boto3): + """If the caller passes just the trailing job id (also valid for + `GetModelInvocationJob`), the logged URL should use it as-is.""" + _, _ = patched_boto3 + logging_obj = MagicMock() + + BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=JOB_ID, aws_region_name="us-west-2", logging_obj=logging_obj + ) + + pre_kwargs = logging_obj.pre_call.call_args.kwargs + assert pre_kwargs["additional_args"]["api_base"] == ( + f"https://bedrock.us-west-2.amazonaws.com/model-invocation-job/{JOB_ID}" + )