scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

feat(bedrock): support retrieve for model-invocation-job batch ARNs

Browse source 
`bedrock.retrieve_batch` previously only handled `:async-invoke/` ARNs
(Twelve Labs Marengo embeddings). The `:model-invocation-job/` ARNs
returned by `CreateModelInvocationJob` (the bulk batch inference API
behind `bedrock.create_batch`) fell through and returned a misleading
data-plane error, leaving created jobs unretrievable through the
LiteLLM batches API.

The two ARN families live on different AWS service endpoints
(`bedrock-runtime` data plane vs `bedrock` control plane), so they need
distinct handlers. This adds:

* `BedrockBatchesHandler._handle_model_invocation_job_status` — calls
  the control plane via boto3 (`bedrock:GetModelInvocationJob`),
  reusing `BaseAWSLLM.get_credentials` for credential resolution so
  model_list / env / role-assumption configs continue to apply. The
  response is reshaped into a `LiteLLMBatch` with the same status
  mapping `transform_create_batch_response` already uses.

* Output-file-URI prediction. Bedrock surfaces the user-supplied
  `s3OutputDataConfig.s3Uri` *prefix* in `GetModelInvocationJob`, but
  results actually land at `<prefix>/<job-id>/<basename(input)>.out`.
  We compute that single-file URI client-side and surface it as
  `output_file_id`, so OpenAI-style `client.files.content(...)` works
  without an extra `ListObjectsV2` round-trip. The bare prefix stays
  in metadata for callers that want the manifest.

* Dispatch in `litellm/batches/main.py` for the new ARN family,
  alongside the existing async-invoke branch.

* Unit tests covering ARN parsing, output-URI prediction (incl. edge
  cases), the full status mapping, region resolution precedence, and
  failure-message propagation.

Note: `request_counts` is intentionally `(0, 0, 0)` —
`GetModelInvocationJob` does not report per-record counts; getting
accurate numbers requires parsing `manifest.json.out` from the output
S3 prefix, which is left to callers.

Made-with: Cursor
This commit is contained in:
Dawei Gu 2026-04-29 16:43:00 -07:00
parent b318231fe9
commit 81e6348b5e
3 changed files with 438 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
litellm/batches/main.py
View file

@ -617,24 +617,35 @@ def retrieve_batch(
_is_async = kwargs.pop("aretrieve_batch", False) is True
client = kwargs.get("client", None)
# Check if this is an async invoke ARN (different from regular batch ARN)
# Async invoke ARNs have format: arn:aws(-[^:]+)?:bedrock:[a-z0-9-]{1,20}:[0-9]{12}:async-invoke/[a-z0-9]{12}
if (
batch_id.startswith("arn:aws")
and ":bedrock:" in batch_id
and ":async-invoke/" in batch_id
):
# Handle async invoke status check
# Remove aws_region_name from kwargs to avoid duplicate parameter
async_kwargs = kwargs.copy()
async_kwargs.pop("aws_region_name", None)
# Bedrock has two distinct ARN families that need different APIs:
# * async-invoke ARNs (Twelve Labs Marengo embeddings) -> bedrock-runtime data plane
# * model-invocation-job ARNs (CreateModelInvocationJob batch) -> bedrock control plane
# They live on different AWS service endpoints and can't share a handler.
# ARN shapes:
# arn:aws(-[^:]+)?:bedrock:<region>:<account>:async-invoke/<id>
# arn:aws(-[^:]+)?:bedrock:<region>:<account>:model-invocation-job/<id>
if batch_id.startswith("arn:aws") and ":bedrock:" in batch_id:
if ":async-invoke/" in batch_id:
# Remove aws_region_name from kwargs to avoid duplicate parameter
async_kwargs = kwargs.copy()
async_kwargs.pop("aws_region_name", None)
return BedrockBatchesHandler._handle_async_invoke_status(
batch_id=batch_id,
aws_region_name=kwargs.get("aws_region_name", "us-east-1"),
logging_obj=litellm_logging_obj,
**async_kwargs,
)
return BedrockBatchesHandler._handle_async_invoke_status(
batch_id=batch_id,
aws_region_name=kwargs.get("aws_region_name", "us-east-1"),
logging_obj=litellm_logging_obj,
**async_kwargs,
)
if ":model-invocation-job/" in batch_id:
mij_kwargs = kwargs.copy()
mij_kwargs.pop("aws_region_name", None)
return BedrockBatchesHandler._handle_model_invocation_job_status(
batch_id=batch_id,
aws_region_name=kwargs.get("aws_region_name"),
logging_obj=litellm_logging_obj,
**mij_kwargs,
)
# Try to use provider config first (for providers like bedrock)
model: Optional[str] = kwargs.get("model", None)

229
litellm/llms/bedrock/batches/handler.py
View file

@ -1,8 +1,79 @@
from datetime import datetime
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
# 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 = {
"Submitted": "validating",
"Validating": "validating",
"Scheduled": "validating",
"InProgress": "in_progress",
"Stopping": "cancelling",
"Stopped": "cancelled",
"Completed": "completed",
"PartiallyCompleted": "completed",
"Failed": "failed",
"Expired": "expired",
}
def _extract_region_from_bedrock_arn(arn: str) -> Optional[str]:
"""ARN shape: ``arn:aws:bedrock:<region>:<account>:<type>/<id>``"""
try:
parts = arn.split(":")
if len(parts) >= 4 and parts[2] == "bedrock":
return parts[3] or None
except Exception:
pass
return None
def _extract_job_id_from_arn(arn: str) -> Optional[str]:
"""``arn:aws:bedrock:<region>:<acct>:model-invocation-job/<job-id>`` -> ``<job-id>``."""
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]:
"""
Compute the deterministic per-job result file URI Bedrock writes to.
Bedrock lays results out as::
<output_prefix>/<job-id>/<basename(input_uri)>.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("/"):
output_prefix = output_prefix + "/"
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) -> Optional[int]:
if value is None:
return None
if isinstance(value, (int, float)):
return int(value)
if isinstance(value, datetime):
return int(value.timestamp())
return None
class BedrockBatchesHandler:
"""
@ -97,3 +168,161 @@ class BedrockBatchesHandler:
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":
"""
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.<region>.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:<region>:<acct>:model-invocation-job/<id>``
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:
raise ImportError(
"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(
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,
)
if logging_obj is not None:
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/"
f"model-invocation-job/{batch_id}"
),
},
)
response = client.get_model_invocation_job(jobIdentifier=batch_id)
if logging_obj is not None:
logging_obj.post_call(
input=batch_id,
api_key="",
original_response=response,
additional_args={"complete_input_dict": {"jobIdentifier": batch_id}},
)
bedrock_status = str(response.get("status", ""))
openai_status = cast(
Any,
_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", "")
)
# Bedrock returns the output *prefix* the user supplied at job creation.
# Actual results land at <prefix>/<job-id>/<basename(input)>.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.
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
)
completed_at = _to_epoch(response.get("endTime"))
openai_batch_metadata: OpenAIBatchMetadata = {
"model_arn": response.get("modelId", ""),
"job_arn": job_arn,
"job_name": response.get("jobName", ""),
"failure_message": response.get("message") or "",
"input_s3_uri": input_uri,
"output_s3_uri": output_prefix,
"output_file_uri": output_file_uri,
}
return LiteLLMBatch(
id=job_arn,
object="batch",
status=openai_status,
created_at=_to_epoch(response.get("submitTime")) or 0,
in_progress_at=_to_epoch(response.get("lastModifiedTime")),
completed_at=completed_at if openai_status == "completed" else None,
failed_at=completed_at if openai_status == "failed" else None,
cancelled_at=completed_at if openai_status == "cancelled" else None,
expired_at=completed_at if openai_status == "expired" else None,
request_counts=BatchRequestCounts(total=0, completed=0, failed=0),
metadata=openai_batch_metadata,
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id=input_uri,
output_file_id=output_file_uri if openai_status == "completed" else None,
)

181
tests/test_litellm/llms/bedrock/batches/test_handler.py Normal file
View file

@ -0,0 +1,181 @@
"""Unit tests for ``BedrockBatchesHandler._handle_model_invocation_job_status``.
These cover the upstream support for retrieving Bedrock bulk batch jobs
(``arn:aws:bedrock:<region>:<acct>:model-invocation-job/<id>``) — the ARN
type returned by ``CreateModelInvocationJob``. We mock the boto3 client so
the tests don't hit AWS.
"""
from __future__ import annotations
import os
import sys
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../../../../.."))
from litellm.llms.bedrock.batches.handler import ( # noqa: E402
BedrockBatchesHandler,
_extract_job_id_from_arn,
_extract_region_from_bedrock_arn,
_predict_output_file_uri,
)
JOB_ID = "abc1234567"
JOB_ARN = f"arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/{JOB_ID}"
INPUT_URI = "s3://my-bucket/inputs/qwen3-235b-a22b-2507-batch.jsonl"
OUTPUT_PREFIX = "s3://my-bucket/litellm-batch-outputs/litellm-bedrock-files-qwen-uuid/"
SUBMIT_TIME = datetime(2026, 4, 28, 12, 0, 0, tzinfo=timezone.utc)
END_TIME = datetime(2026, 4, 28, 12, 30, 0, tzinfo=timezone.utc)
def _fake_boto3_response(status: str = "Completed", end_time=END_TIME):
return {
"jobArn": JOB_ARN,
"jobName": "litellm-bedrock-files-qwen-uuid",
"modelId": "bedrock/qwen.qwen3-235b-a22b-2507-v1:0",
"status": status,
"submitTime": SUBMIT_TIME,
"lastModifiedTime": end_time,
"endTime": end_time,
"inputDataConfig": {"s3InputDataConfig": {"s3Uri": INPUT_URI}},
"outputDataConfig": {"s3OutputDataConfig": {"s3Uri": OUTPUT_PREFIX}},
}
@pytest.fixture
def patched_boto3():
"""Yield a stub bedrock client whose `get_model_invocation_job` is a MagicMock."""
fake_client = MagicMock()
fake_client.get_model_invocation_job.return_value = _fake_boto3_response()
with (
patch("boto3.client", return_value=fake_client) as boto_client_factory,
patch(
"litellm.llms.bedrock.batches.transformation.BedrockBatchesConfig.get_credentials",
return_value=MagicMock(access_key="AKIA", secret_key="SECRET", token=None),
),
):
yield fake_client, boto_client_factory
def test_extract_region_from_arn():
assert _extract_region_from_bedrock_arn(JOB_ARN) == "us-west-2"
assert _extract_region_from_bedrock_arn("arn:aws:bedrock::123:foo/bar") is None
assert _extract_region_from_bedrock_arn("not-an-arn") is None
def test_extract_job_id_from_arn():
assert _extract_job_id_from_arn(JOB_ARN) == JOB_ID
assert (
_extract_job_id_from_arn("arn:aws:bedrock:us-west-2:1:async-invoke/x") is None
)
def test_predict_output_file_uri_happy_path():
expected = f"{OUTPUT_PREFIX}{JOB_ID}/qwen3-235b-a22b-2507-batch.jsonl.out"
assert _predict_output_file_uri(OUTPUT_PREFIX, INPUT_URI, JOB_ID) == expected
def test_predict_output_file_uri_adds_trailing_slash():
prefix_no_slash = OUTPUT_PREFIX.rstrip("/")
expected = f"{OUTPUT_PREFIX}{JOB_ID}/qwen3-235b-a22b-2507-batch.jsonl.out"
assert _predict_output_file_uri(prefix_no_slash, INPUT_URI, JOB_ID) == expected
@pytest.mark.parametrize(
"missing_arg",
[
("", INPUT_URI, JOB_ID),
(OUTPUT_PREFIX, "", JOB_ID),
(OUTPUT_PREFIX, INPUT_URI, None),
],
)
def test_predict_output_file_uri_returns_none_when_missing_input(missing_arg):
assert _predict_output_file_uri(*missing_arg) is None
def test_handle_model_invocation_job_status_completed(patched_boto3):
fake_client, boto_client_factory = patched_boto3
batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
fake_client.get_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN)
# Region should be sniffed from the ARN.
_, kwargs = boto_client_factory.call_args
assert kwargs["region_name"] == "us-west-2"
assert batch.id == JOB_ARN
assert batch.status == "completed"
assert batch.input_file_id == INPUT_URI
expected_out = f"{OUTPUT_PREFIX}{JOB_ID}/qwen3-235b-a22b-2507-batch.jsonl.out"
assert batch.output_file_id == expected_out
assert batch.completed_at == int(END_TIME.timestamp())
assert batch.failed_at is None
assert batch.cancelled_at is None
# Per-record counts aren't reported by GetModelInvocationJob, so we leave
# them zeroed; consumers should parse manifest.json.out for accurate counts.
assert batch.request_counts.total == 0
assert batch.metadata["job_arn"] == JOB_ARN
assert batch.metadata["output_file_uri"] == expected_out
assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX
@pytest.mark.parametrize(
"bedrock_status,openai_status",
[
("Submitted", "validating"),
("Validating", "validating"),
("Scheduled", "validating"),
("InProgress", "in_progress"),
("Stopping", "cancelling"),
("Stopped", "cancelled"),
("Completed", "completed"),
("PartiallyCompleted", "completed"),
("Failed", "failed"),
("Expired", "expired"),
# Unknown/unmapped Bedrock status falls back to "in_progress" so we
# don't 500 on a future AWS-side enum addition.
("MyBrandNewStatus", "in_progress"),
],
)
def test_status_mapping(patched_boto3, bedrock_status, openai_status):
fake_client, _ = patched_boto3
fake_client.get_model_invocation_job.return_value = _fake_boto3_response(
status=bedrock_status
)
batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
assert batch.status == openai_status
# output_file_id is only populated for terminal-completed jobs, so callers
# don't accidentally try to download a non-existent file mid-run.
if openai_status == "completed":
assert batch.output_file_id is not None
else:
assert batch.output_file_id is None
def test_explicit_region_overrides_arn(patched_boto3):
_, boto_client_factory = patched_boto3
BedrockBatchesHandler._handle_model_invocation_job_status(
batch_id=JOB_ARN, aws_region_name="eu-central-1"
)
_, kwargs = boto_client_factory.call_args
assert kwargs["region_name"] == "eu-central-1"
def test_failure_message_propagates(patched_boto3):
fake_client, _ = patched_boto3
failed_response = _fake_boto3_response(status="Failed")
failed_response["message"] = "Input file failed validation"
fake_client.get_model_invocation_job.return_value = failed_response
batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
assert batch.status == "failed"
assert batch.failed_at == int(END_TIME.timestamp())
assert batch.metadata["failure_message"] == "Input file failed validation"