Merge pull request #39847 from BerriAI/litellm_lit_5730_bedrock_batch_cancel_e2e

test(e2e/batches): assert Bedrock batch cancel and list in the lifecycle

tests/e2e/batches/COVERAGE.md

@@ -19,11 +19,13 @@ failures are hard test failures (see `tests/e2e/CLAUDE.md`).
 | OpenAI    | yes | yes | yes | yes | yes (lifecycle + terminal output) | OpenAI Files |
 | Azure     | yes | yes | yes | yes | yes (byte-verbatim) | Azure Files |
 | Vertex AI | yes | yes | yes | yes | yes (provider-transformed) | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) |
-| Bedrock   | yes (unified only) | yes | no (limited upstream) | no | yes (provider-transformed) | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) |
+| Bedrock   | yes (unified only) | yes | yes | yes (unfiltered managed list) | yes (provider-transformed) | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) |
 
-Bedrock cancel is unreliable upstream and list is unsupported, so both are gated off
-(`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix;
-flipping those gates is tracked in LIT-4774 and deliberately not part of this suite.
+Bedrock cancel maps to `StopModelInvocationJob` and comes back `cancelling`; the
+lifecycle asserts it the same way it does for OpenAI (`_CANCEL_ASSERTED_PROVIDERS`).
+Bedrock has no provider-side list, so list is the proxy's DB-backed managed view: the
+unified lifecycle lists with the plain `GET /v1/batches` and the batch must appear
+there. Both were gated off until LIT-5730, after LIT-4774 landed cancel support. A batch that completes inside the 2 s pre-cancel window skips the cancel assertion (a documented vacuous pass for the cancel cell, same as OpenAI); the list assertion runs either way.
 Bedrock file upload requires a model on the request (`encoded` / `unified` scenarios only);
 `model_param` and `provider_fallback` are omitted because `POST /bedrock/v1/files` has no
 model-less passthrough path.
@@ -148,6 +150,6 @@ never landed.
 Unified (managed) batch cost is owned by the hourly `CheckBatchCost` poller, and a
 terminal DB status short-circuits retrieve for those ids, so the terminal-state cell
 uses the encoded path; poller timing does not fit an e2e gate and belongs in a
-DI-stubbed proxy integration test under `tests/test_litellm/proxy/`. Bedrock
-cancel/list stay gated pending LIT-4774. Gemini (non-Vertex) file content raises
-`NotImplementedError` upstream and is not a coverage cell.
+DI-stubbed proxy integration test under `tests/test_litellm/proxy/`. Gemini
+(non-Vertex) file content raises `NotImplementedError` upstream and is not a
+coverage cell.

tests/e2e/batches/capabilities.py

@@ -143,8 +143,8 @@ PROVIDERS: tuple[Provider, ...] = (
         "bedrock",
         batch_model_name("bedrock-batch"),
         "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
-        can_cancel=False,
-        can_list=False,
+        can_cancel=True,
+        can_list=True,
     ),
 )
 
@@ -248,8 +248,9 @@ def coverage_cells_for_lifecycle(cap: Capability) -> tuple[str, ...]:
     """Registry cell ids that the parametrized lifecycle test covers for one capability.
 
     OpenAI has per-scenario cells plus granular create/retrieve/cancel/list/file
-    cells. Other providers have one basic cell each. File-upload cells for the
-    batch-backing path are included when the lifecycle uploads for that provider.
+    cells. Bedrock adds cancel and list cells behind its gates. Other providers
+    have one basic cell each. File-upload cells for the batch-backing path are
+    included when the lifecycle uploads for that provider.
     """
     match cap.provider:
         case "openai":
@@ -279,6 +280,8 @@ def coverage_cells_for_lifecycle(cap: Capability) -> tuple[str, ...]:
             return (
                 "llm.batches.bedrock.basic.nonstream.works",
                 "llm.files.bedrock.upload.nonstream.works",
+                *(("llm.batches.bedrock.cancel.nonstream.works",) if cap.can_cancel else ()),
+                *(("llm.batches.bedrock.list.nonstream.works",) if cap.can_list else ()),
             )
         case _:
             return ()

tests/e2e/batches/test_batches_e2e.py

@@ -77,7 +77,7 @@ BATCH_OP_RETRIES = 5
 # (connection refused, brief 500s) and the registry only has one basic cell per
 # provider (shared across scenarios). Create + retrieve already prove routing;
 # cancel is still deferred for cleanup, just not asserted for these two.
-_CANCEL_ASSERTED_PROVIDERS = frozenset({"openai"})
+_CANCEL_ASSERTED_PROVIDERS = frozenset({"openai", "bedrock"})
 
 
 def _transient_status(status_code: int) -> bool:
@@ -293,14 +293,13 @@ def test_batch_lifecycle(
             f"batch reached {pre_cancel.status!r} before cancel; "
             "provider likely rejected the input"
         )
-        if pre_cancel.status == "completed":
-            return
-        cancelled = cancel_batch(client, batch.id, key=key, provider=provider)
-        assert cancelled.id == batch.id
-        assert cancelled.object == "batch"
-        assert cancelled.status in {"cancelling", "cancelled"}, (
-            f"unexpected post-cancel status {cancelled.status!r}"
-        )
+        if pre_cancel.status != "completed":
+            cancelled = cancel_batch(client, batch.id, key=key, provider=provider)
+            assert cancelled.id == batch.id
+            assert cancelled.object == "batch"
+            assert cancelled.status in {"cancelling", "cancelled"}, (
+                f"unexpected post-cancel status {cancelled.status!r}"
+            )
 
     if cap.can_list:
         list_result = client.list_batches(key=key, provider=provider)

tests/e2e/coverage_registry/llm_nonconversational.yaml

@@ -23,6 +23,8 @@
 - {id: llm.batches.vertex.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Vertex batches"}
 - {id: llm.batches.bedrock.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Bedrock batches (encoded/unified only)"}
 - {id: llm.batches.bedrock.assume_role.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: assume_role, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock batch create under STS assume-role credentials"}
+- {id: llm.batches.bedrock.cancel.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock batch cancel (StopModelInvocationJob) returns the same id with a cancelling/cancelled status"}
+- {id: llm.batches.bedrock.list.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "A Bedrock managed batch is present in the GET /v1/batches list envelope"}
 - {id: llm.batches.hosted_vllm.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible batch create"}
 - {id: llm.batches.openai.key_model_access_denied.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Key model restriction 403 on upload/create"}
 - {id: llm.batches.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.18 / LIT-4778", rationale: "Missing input_file_id and invalid batch id rejected"}