diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 919c39f21a2..b36d8937ad0 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/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`). diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index c4b699190b8..eff8f297f25 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/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" diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 635ea3f7ea5..50f9b9808b2 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/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)"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 03d15f532b8..fa6dad90126 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/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",