Merge pull request #41536 from BerriAI/litellm_aws_govcloud_partition_gate

test(e2e): cover bedrock batch file upload and create in the us-gov-west-1 partition

tests/e2e/batches/COVERAGE.md

@@ -20,6 +20,7 @@ failures are hard test failures (see `tests/e2e/CLAUDE.md`).
 | 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 | yes | yes (unfiltered managed list) | yes (provider-transformed) | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) |
+| Bedrock GovCloud (`us-gov-west-1`) | yes (unified only) | yes | no | no | yes (provider-transformed) | S3 (`s3_bucket_name` + `aws_*` on model, resolved from `AWS_GOVCLOUD_ACCESS_KEY_ID` / `AWS_GOVCLOUD_SECRET_ACCESS_KEY` / `AWS_GOVCLOUD_BATCH_S3_BUCKET` / `AWS_GOVCLOUD_BATCH_ROLE_ARN`) |
 
 Bedrock cancel maps to `StopModelInvocationJob` and comes back `cancelling`; the
 lifecycle asserts it the same way it does for OpenAI (`_CANCEL_ASSERTED_PROVIDERS`).

tests/e2e/batches/test_batches_e2e.py

@@ -21,21 +21,18 @@ import os
 import re
 import time
 from datetime import datetime, timedelta, timezone
+from typing import Final
 
 import pytest
-from pydantic import BaseModel
-
-from e2e_config import MASTER_KEY, PROXY_BASE_URL, unique_marker
-
 from batch_cleanup import cleanup_batch, cleanup_file
 from batch_client import (
     AZURE_FILE_EXPIRY_SECONDS,
-    batch_upload_form,
     UPLOAD_FILENAME,
     BatchClient,
     BatchCreateBody,
     BatchObject,
     FileObject,
+    batch_upload_form,
     is_model_access_denied,
     is_result_access_denied,
 )
@@ -57,6 +54,7 @@ from capabilities import (
     openai_batch_params,
     raw_id_matches_provider,
 )
+from e2e_config import MASTER_KEY, PROXY_BASE_URL, unique_marker
 from e2e_http import (
     FileUploadForm,
     Result,
@@ -68,6 +66,7 @@ from e2e_http import (
 )
 from lifecycle import ResourceManager
 from models import KeyGenerateBody, KeyMetadata, LiteLLMParamsBody, SpendLogRow
+from pydantic import BaseModel, Field
 
 pytestmark = pytest.mark.e2e
 
@@ -75,6 +74,25 @@ CREATED_BATCH_STATUSES = {"validating", "in_progress", "finalizing"}
 BATCH_CANCEL_DELAY_SECONDS = 2
 BATCH_TERMINAL_BEFORE_CANCEL = {"failed", "cancelled", "expired"}
 BATCH_OP_RETRIES = 5
+
+
+class _GovCloudBedrockContent(BaseModel):
+    text: str
+
+
+class _GovCloudBedrockMessage(BaseModel):
+    content: tuple[_GovCloudBedrockContent, ...]
+
+
+class _GovCloudBedrockInput(BaseModel):
+    messages: tuple[_GovCloudBedrockMessage, ...]
+
+
+class _GovCloudBedrockRecord(BaseModel):
+    record_id: str = Field(alias="recordId")
+    model_input: _GovCloudBedrockInput = Field(alias="modelInput")
+
+
 # Azure / Vertex cancel and the pre-cancel re-retrieve are provider-side flakes
 # (connection refused, brief 500s) and the registry only has one basic cell per
 # provider (shared across scenarios). Create + retrieve already prove routing;
@@ -1006,6 +1024,91 @@ class TestBedrockBatchAssumeRole:
         assert fetched.id == batch.id
 
 
+GOVCLOUD_REGION: Final = "us-gov-west-1"
+GOVCLOUD_RAW_MODEL: Final = "bedrock/amazon.nova-lite-v1:0"
+
+
+def _govcloud_params() -> LiteLLMParamsBody:
+    return LiteLLMParamsBody(
+        model=GOVCLOUD_RAW_MODEL,
+        aws_access_key_id="os.environ/AWS_GOVCLOUD_ACCESS_KEY_ID",
+        aws_secret_access_key="os.environ/AWS_GOVCLOUD_SECRET_ACCESS_KEY",
+        aws_region_name=GOVCLOUD_REGION,
+        s3_region_name=GOVCLOUD_REGION,
+        s3_bucket_name="os.environ/AWS_GOVCLOUD_BATCH_S3_BUCKET",
+        s3_access_key_id="os.environ/AWS_GOVCLOUD_ACCESS_KEY_ID",
+        s3_secret_access_key="os.environ/AWS_GOVCLOUD_SECRET_ACCESS_KEY",
+        aws_batch_role_arn="os.environ/AWS_GOVCLOUD_BATCH_ROLE_ARN",
+    )
+
+
+class TestBedrockBatchGovCloud:
+    """Bedrock batch lifecycle in the AWS GovCloud partition (us-gov-west-1).
+
+    The deployment carries a GovCloud region for both Bedrock and S3, so the proxy has to
+    sign the file upload against the us-gov S3 endpoint and submit the job to the us-gov
+    Bedrock endpoint. Commercial-partition hostnames or arn:aws: ARNs reject the GovCloud
+    key, so a partition regression fails the upload instead of passing silently.
+    """
+
+    @pytest.mark.covers(
+        "llm.batches.bedrock.govcloud_partition.nonstream.works",
+        "llm.files.bedrock.govcloud_partition.nonstream.works",
+        exercised_on=["batches", "files"],
+    )
+    def test_unified_file_upload_and_batch_create_in_govcloud(
+        self, client: BatchClient, resources: ResourceManager
+    ) -> None:
+        model_name: Final = batch_model_name("bedrock-govcloud-batch")
+        model_id: Final = client.create_model(model_name, _govcloud_params())
+        resources.defer(lambda: client.delete_model(model_id))
+        key: Final = resources.key()
+        file: Final = unwrap(
+            client.upload_file(
+                content=render_jsonl(GOVCLOUD_RAW_MODEL),
+                form=FileUploadForm(purpose="batch", target_model_names=model_name),
+                key=key,
+            )
+        )
+        resources.defer(lambda: cleanup_file(client, file.id, key=key))
+        assert_file_object(file, provider="bedrock")
+
+        downloaded: Final = client.proxy.transport.download(
+            f"/v1/files/{file.id}/content",
+            headers=client.proxy.transport.bearer(key),
+        )
+        assert downloaded.status_code == 200, (
+            f"GovCloud file content must be 200, got {downloaded.status_code}: {downloaded.body[:300]}"
+        )
+        downloaded_lines: Final = downloaded.body.strip().splitlines()
+        assert len(downloaded_lines) == 1, (
+            f"GovCloud file content download must contain one JSONL record, got {len(downloaded_lines)}"
+        )
+        downloaded_record: Final = _GovCloudBedrockRecord.model_validate(json.loads(downloaded_lines[0]))
+        assert downloaded_record.record_id == "req-1", (
+            f"GovCloud file content must preserve the uploaded custom_id, got {downloaded_record.record_id!r}"
+        )
+        assert downloaded_record.model_input.messages[0].content[0].text == "ping", (
+            "GovCloud file content must preserve the uploaded message text"
+        )
+
+        created: Final = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
+        require_successful_call(created)
+        batch: Final = BatchObject.model_validate_json(created.body)
+        resources.defer(lambda: cleanup_batch(client, batch.id, key=key))
+
+        assert is_managed_id(batch.id), (
+            f"GovCloud create via target_model_names must return a managed batch id, got {batch.id!r}"
+        )
+        assert batch.status in CREATED_BATCH_STATUSES, (
+            f"GovCloud batch has non-transitional status {batch.status!r}"
+        )
+        assert_batch_object(batch)
+
+        fetched: Final = unwrap(client.retrieve_batch(batch.id, key=key))
+        assert fetched.id == batch.id
+
+
 GEMINI_FILES_RAW_MODEL = "gemini-2.5-flash"
 
 

tests/e2e/coverage_registry/llm_nonconversational.yaml

@@ -23,6 +23,7 @@
 - {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.govcloud_partition.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: govcloud_partition, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock batch create in the us-gov-west-1 partition"}
 - {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"}
@@ -45,6 +46,7 @@
 - {id: llm.files.azure_openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:45", rationale: "Azure file upload managed backend"}
 - {id: llm.files.vertex.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:52", rationale: "Vertex file upload to GCS"}
 - {id: llm.files.bedrock.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:59", rationale: "Bedrock file upload to S3"}
+- {id: llm.files.bedrock.govcloud_partition.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: govcloud_partition, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock file upload to an S3 bucket in the us-gov-west-1 partition"}
 - {id: llm.files.gemini.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: gemini, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Gemini Files API upload via proxy"}
 - {id: llm.files.hosted_vllm.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible file upload"}
 - {id: llm.files.openai.require_managed_files_upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_managed_files_enforcement_e2e.py / LIT-5902", rationale: "With require_managed_files enabled, an upload without target_model_names and an upload carrying a model param are both rejected 400; runs only in the sequential managed-files stack phase (E2E_MANAGED_FILES_STACK)"}

tests/e2e/coverage_registry/schema.py

@@ -64,6 +64,7 @@ LlmCapability = Literal[
     "assume_role",
     "basic",
     "count_tokens",
+    "govcloud_partition",
     "input_validation",
     "long_context_1m",
     "mid_conversation_system",