From e53e67ede556d3d41d23b35c18f5dc16dda7671a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 26 Sep 2026 11:17:03 -0700 Subject: [PATCH] test(e2e): assert only litellm-owned batch behavior and move the blank S3 env pin to an integration test (#43321) --- tests/e2e/batches/COVERAGE.md | 10 +- tests/e2e/batches/batch_cleanup.py | 30 ++- tests/e2e/batches/bedrock_env_gateway.py | 151 ------------ tests/e2e/batches/test_batch_cleanup.py | 71 +++++- tests/e2e/batches/test_batches_e2e.py | 6 +- .../batches/test_bedrock_blank_s3_env_e2e.py | 109 --------- .../llm_nonconversational.yaml | 1 - tests/e2e/coverage_registry/schema.py | 1 - .../test_bedrock_batch_blank_s3_env_wire.py | 222 ++++++++++++++++++ 9 files changed, 319 insertions(+), 282 deletions(-) delete mode 100644 tests/e2e/batches/bedrock_env_gateway.py delete mode 100644 tests/e2e/batches/test_bedrock_blank_s3_env_e2e.py create mode 100644 tests/integration/providers/test_bedrock_batch_blank_s3_env_wire.py diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 862eef5c0f4..cd0fb35165e 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -22,10 +22,9 @@ failures are hard test failures (see `tests/e2e/AGENTS.md`). | 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 split S3 identity | no | no | no | no | yes (file upload, content, delete) | S3 signed with `s3_access_key_id` / `s3_secret_access_key` (`AWS_S3_ONLY_ACCESS_KEY_ID` / `AWS_S3_ONLY_SECRET_ACCESS_KEY`, object rights on `AWS_BATCH_S3_BUCKET` only) while `aws_*` is `AWS_BEDROCK_ONLY_ACCESS_KEY_ID` / `AWS_BEDROCK_ONLY_SECRET_ACCESS_KEY`, an identity with no S3 rights on that bucket | -| Bedrock blank S3 env | yes (unified only, on an owned gateway exporting `AWS_S3_ENCRYPTION_KEY_ID` / `AWS_S3_BUCKET_OWNER` as empty strings) | no | no | no | no | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` in the gateway config); blank env vars must be treated as unset, not serialized | Bedrock cancel maps to `StopModelInvocationJob` and comes back `cancelling`; the -lifecycle asserts it the same way it does for OpenAI (`_CANCEL_ASSERTED_PROVIDERS`). +lifecycle asserts it the same way it does for OpenAI and Azure (`_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. @@ -132,8 +131,11 @@ provider when deleted. Model-encoded and managed file IDs route themselves File deletion and batch cancellation check their responses and retry transient failures up to three times. Teardown attempts every registered cleanup before reporting failures as test errors. Already deleted files and batches that are -terminal are safe to clean up again. Managed batch cancellation polls for up to eleven minutes -before input deletion: the ten-minute provider window plus a propagation margin. +terminal are safe to clean up again. Managed batch cancellation polls for up to two minutes +before input deletion. A managed batch still `cancelling` after that is left for the provider to +finish, and its input file is left in place because LiteLLM refuses to delete a file a non-terminal +batch references. Both are reported as `BatchCleanupLeftover` warnings naming their ids rather than +failing the test. Any other status or error still fails Accepted cancellation may still report validating or in_progress while the provider updates its state. Raw and model-encoded batches are polled until cancelling or terminal before input deletion. OpenAI and Azure lifecycle cleanup also deletes diff --git a/tests/e2e/batches/batch_cleanup.py b/tests/e2e/batches/batch_cleanup.py index f889844f1ae..5b3baaa624c 100644 --- a/tests/e2e/batches/batch_cleanup.py +++ b/tests/e2e/batches/batch_cleanup.py @@ -1,3 +1,4 @@ +import warnings from builtins import ExceptionGroup from collections.abc import Callable from itertools import count @@ -12,8 +13,9 @@ from pydantic import BaseModel CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0) BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"}) BATCH_PENDING_STATUSES: Final = frozenset({"validating", "in_progress", "finalizing", "cancelling"}) -BATCH_CANCEL_TIMEOUT_SECONDS: Final = 660.0 +BATCH_CANCEL_TIMEOUT_SECONDS: Final = 120.0 BATCH_CANCEL_POLL_SECONDS: Final = 10.0 +FILE_IN_USE_REFUSAL: Final = "batch(es) in non-terminal state" class BatchCleanupClient(Protocol): @@ -26,6 +28,10 @@ class BatchCleanupClient(Protocol): def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ... +class BatchCleanupLeftover(UserWarning): + pass + + def cleanup_result[R: BaseModel]( action: Callable[[], Result[R]], *, wait: Callable[[float], None] = sleep ) -> Result[R]: @@ -59,6 +65,13 @@ def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider result: Final = cleanup_result(delete) if isinstance(result, UnknownApiError) and result.status_code == 404: return + if isinstance(result, UnknownApiError) and result.status_code == 400 and FILE_IN_USE_REFUSAL in result.body: + warnings.warn( + f"Left file {file_id} in place: LiteLLM refused to delete it while a batch still references it", + BatchCleanupLeftover, + stacklevel=2, + ) + return deleted: Final = _require_cleanup_success(result, f"Delete file {file_id}") assert deleted.deleted is True or ( deleted.deleted is None and is_managed_id(file_id) and deleted.id == file_id and deleted.object == "file" @@ -120,10 +133,17 @@ def cleanup_batch( ) if current.status == "cancelling" and not needs_terminal_state: return - assert clock() < deadline, ( - f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s, " - f"last status {current.status}" - ) + if clock() >= deadline: + assert current.status == "cancelling", ( + f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s, " + f"last status {current.status}" + ) + warnings.warn( + f"Left batch {batch_id} cancelling after {BATCH_CANCEL_TIMEOUT_SECONDS}s for the provider to finish", + BatchCleanupLeftover, + stacklevel=2, + ) + return wait(BATCH_CANCEL_POLL_SECONDS) diff --git a/tests/e2e/batches/bedrock_env_gateway.py b/tests/e2e/batches/bedrock_env_gateway.py deleted file mode 100644 index 0b1d840eb30..00000000000 --- a/tests/e2e/batches/bedrock_env_gateway.py +++ /dev/null @@ -1,151 +0,0 @@ -"""An owned, source-built proxy whose process env exports AWS_S3_* vars blank. - -The shared fixture proxy inherits the harness env, which cannot reproduce a user -shell that exports AWS_S3_ENCRYPTION_KEY_ID / AWS_S3_BUCKET_OWNER as empty -strings. This gateway boots a second proxy with both vars present but blank, so -a batch create through it proves blank means unset, not an empty string. -""" - -from __future__ import annotations - -import importlib.util -import os -import shutil -import socket -import subprocess -import sys -import tempfile -import time -from collections.abc import Mapping -from dataclasses import dataclass, field -from pathlib import Path -from typing import Final - -from e2e_config import unique_marker -from e2e_http import NoBody -from idp import stop_process_group -from proxy_client import ProxyClient, build_proxy_client -from pydantic import TypeAdapter - -STARTUP_TIMEOUT_SECONDS: Final = 240 -LOG_TAIL_BYTES: Final = 4000 - - -def litellm_root() -> Path: - spec: Final = importlib.util.find_spec("litellm") - assert spec is not None and spec.origin is not None, "litellm must be importable to boot the blank-S3-env gateway" - return Path(spec.origin).resolve().parents[1] - - -_CONFIG_YAML: Final = """model_list: - - model_name: bedrock-blank-s3-batch - litellm_params: - model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0 - aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID - aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY - aws_region_name: os.environ/AWS_REGION - s3_region_name: os.environ/AWS_REGION - s3_bucket_name: os.environ/AWS_BATCH_S3_BUCKET - s3_access_key_id: os.environ/AWS_ACCESS_KEY_ID - s3_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY - aws_batch_role_arn: os.environ/AWS_BATCH_ROLE_ARN - -general_settings: - master_key: os.environ/LITELLM_MASTER_KEY - database_url: os.environ/DATABASE_URL -""" - - -def available_port() -> int: - with socket.socket() as listener: - listener.bind(("127.0.0.1", 0)) - return TypeAdapter(tuple[str, int]).validate_python(listener.getsockname())[1] - - -@dataclass(slots=True) -class BedrockEnvGateway: - base_url: str - master_key: str - proxy: ProxyClient - _environment: Mapping[str, str] = field(repr=False) - _command: tuple[str, ...] = field(repr=False) - _log_path: Path - _child: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False) - - @classmethod - def start(cls) -> BedrockEnvGateway: - assert os.environ.get("DATABASE_URL"), "DATABASE_URL is required for the blank-S3-env gateway" - root: Final = litellm_root() - port: Final = available_port() - base_url: Final = f"http://127.0.0.1:{port}" - master_key: Final = f"sk-e2e-blank-s3-{unique_marker()}" - directory: Final = Path(tempfile.mkdtemp(prefix="litellm-e2e-blank-s3-")) - config: Final = directory / "blank-s3-gateway.yaml" - config.write_text(_CONFIG_YAML) - environment: Final = { - **{key: value for key, value in os.environ.items() if not key.startswith("REDIS_")}, - "DATABASE_URL": os.environ["DATABASE_URL"], - "LITELLM_MASTER_KEY": master_key, - "STORE_MODEL_IN_DB": "False", - "PYTHONPATH": str(root), - "AWS_S3_ENCRYPTION_KEY_ID": "", - "AWS_S3_BUCKET_OWNER": "", - } - gateway: Final = cls( - base_url=base_url, - master_key=master_key, - proxy=build_proxy_client( - base_url=base_url, - control_plane_base_url=base_url, - replica_urls=(base_url,), - master_key=master_key, - ), - _environment=environment, - _command=( - sys.executable, - "-m", - "litellm.proxy.proxy_cli", - "--config", - str(config), - "--port", - str(port), - "--host", - "127.0.0.1", - ), - _log_path=directory / "blank-s3-gateway.log", - ) - with gateway._log_path.open("ab") as log: - gateway._child = subprocess.Popen( - gateway._command, - env=dict(gateway._environment), - stdout=log, - stderr=log, - start_new_session=True, - cwd=root, - ) - deadline: Final = time.monotonic() + STARTUP_TIMEOUT_SECONDS - while time.monotonic() < deadline: - assert gateway._child.poll() is None, f"blank-S3-env gateway exited early; log tail:\n{gateway.log_tail()}" - result = gateway.proxy.transport.probe("/health/liveliness", params=NoBody()) - if result.status_code == 200: - return gateway - time.sleep(0.5) - tail: Final = gateway.log_tail() - gateway.stop() - raise AssertionError( - f"blank-S3-env gateway did not become ready in {STARTUP_TIMEOUT_SECONDS}s; log tail:\n{tail}" - ) - - def log_tail(self) -> str: - if not self._log_path.exists(): - return "" - with self._log_path.open("rb") as log: - log.seek(0, 2) - size: Final = log.tell() - log.seek(max(0, size - LOG_TAIL_BYTES)) - return log.read().decode("utf-8", errors="replace") - - def stop(self) -> None: - if self._child is not None: - stop_process_group(self._child) - shutil.rmtree(self._log_path.parent, ignore_errors=True) diff --git a/tests/e2e/batches/test_batch_cleanup.py b/tests/e2e/batches/test_batch_cleanup.py index 28eb362e876..5e2ac12d300 100644 --- a/tests/e2e/batches/test_batch_cleanup.py +++ b/tests/e2e/batches/test_batch_cleanup.py @@ -4,7 +4,14 @@ from typing import Final from unittest.mock import Mock, call import pytest -from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result +from batch_cleanup import ( + BATCH_CANCEL_TIMEOUT_SECONDS, + CLEANUP_DELAYS, + BatchCleanupLeftover, + cleanup_batch, + cleanup_file, + cleanup_result, +) from batch_client import AZURE_FILE_EXPIRY_SECONDS, BatchObject, FileDeleteResponse, batch_upload_form from capabilities import CAPABILITIES, Capability from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError @@ -13,6 +20,10 @@ from models import KeyGenerateBody MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE=" MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x" +IN_USE_REFUSAL: Final = ( + f'{{"error":{{"message":"Cannot delete file {MANAGED_FILE_ID}. The file is referenced by 1 batch(es) in ' + f'non-terminal state: {MANAGED_BATCH_ID}: cancelling. ","type":"invalid_request_error","code":"400"}}}}' +) class ExpectedCalls[T]: @@ -125,6 +136,29 @@ class TestFileCleanup: cleanup_file(client, "file-1", key="test-key") client.calls.assert_done() + def test_delete_refused_because_a_batch_still_references_the_file_is_left_and_reported(self) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls((f"delete None {MANAGED_FILE_ID}",)), + files=(UnknownApiError(status_code=400, body=IN_USE_REFUSAL),), + ) + with pytest.warns(BatchCleanupLeftover, match=MANAGED_FILE_ID): + cleanup_file(client, MANAGED_FILE_ID, key="test-key") + client.calls.assert_done() + + @pytest.mark.parametrize( + "failure", + [ + UnknownApiError(status_code=400, body="Invalid file id"), + UnknownApiError(status_code=409, body=IN_USE_REFUSAL), + UnknownApiError(status_code=501, body=IN_USE_REFUSAL), + ], + ) + def test_any_other_delete_failure_still_raises(self, failure: UnknownApiError) -> None: + client: Final = CleanupClient(calls=ExpectedCalls((f"delete None {MANAGED_FILE_ID}",)), files=(failure,)) + with pytest.raises(AssertionError, match=f"Delete file {MANAGED_FILE_ID} failed: HTTP {failure.status_code}"): + cleanup_file(client, MANAGED_FILE_ID, key="test-key") + client.calls.assert_done() + def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None: client: Final = CleanupClient( calls=ExpectedCalls(("delete azure file-1",)), @@ -188,29 +222,50 @@ class TestBatchCancellation: client.calls.assert_done() delays.assert_done() - def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None: + def test_batch_still_cancelling_at_the_deadline_and_its_input_file_are_left_and_reported(self) -> None: client: Final = CleanupClient( calls=ExpectedCalls( ( f"retrieve None {MANAGED_BATCH_ID}", f"retrieve None {MANAGED_BATCH_ID}", - "delete None file-1", + f"delete None {MANAGED_FILE_ID}", "delete key test-key", ) ), batches=(batch("cancelling"), batch("cancelling")), - files=(deleted_file(),), + files=(UnknownApiError(status_code=400, body=IN_USE_REFUSAL),), ) times: Final = (0.0, BATCH_CANCEL_TIMEOUT_SECONDS) ticks: Final[Callable[[], float]] = Mock(side_effect=times) manager: Final = ResourceManager(client=client, strict_cleanup=True) key: Final = manager.key() - manager.defer(lambda: cleanup_file(client, "file-1", key=key)) + manager.defer(lambda: cleanup_file(client, MANAGED_FILE_ID, key=key)) manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=ticks)) - with pytest.raises(ExceptionGroup) as caught: + with pytest.warns(BatchCleanupLeftover) as leftovers: manager.teardown() - assert "cancellation did not finish" in str(caught.value.exceptions[0]) - assert "last status cancelling" in str(caught.value.exceptions[0]) + client.calls.assert_done() + messages: Final = tuple(str(warning.message) for warning in leftovers) + assert len(messages) == 2 + assert MANAGED_BATCH_ID in messages[0] and "cancelling" in messages[0] + assert MANAGED_FILE_ID in messages[1] + + @pytest.mark.parametrize( + "last, reported", + [ + (batch("in_progress"), f"did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s, last status in_progress"), + (UnknownApiError(status_code=403, body="forbidden"), "after cancellation failed: HTTP 403"), + ], + ) + def test_anything_but_still_cancelling_at_the_deadline_still_fails( + self, last: Result[BatchObject], reported: str + ) -> None: + client: Final = CleanupClient( + calls=ExpectedCalls((f"retrieve None {MANAGED_BATCH_ID}",) * 2), batches=(batch("cancelling"), last) + ) + times: Final = (0.0, BATCH_CANCEL_TIMEOUT_SECONDS) + ticks: Final[Callable[[], float]] = Mock(side_effect=times) + with pytest.raises(AssertionError, match=reported): + cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", clock=ticks) client.calls.assert_done() @pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"]) diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 6e9cf45e787..8da2deb4010 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -94,11 +94,11 @@ class _GovCloudBedrockRecord(BaseModel): model_input: _GovCloudBedrockInput = Field(alias="modelInput") -# Azure / Vertex cancel and the pre-cancel re-retrieve are provider-side flakes +# 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; -# cancel is still deferred for cleanup, just not asserted for these two. -_CANCEL_ASSERTED_PROVIDERS = frozenset({"openai", "bedrock"}) +# cancel is still deferred for cleanup, just not asserted for Vertex. +_CANCEL_ASSERTED_PROVIDERS = frozenset({"openai", "azure", "bedrock"}) def _transient_status(status_code: int) -> bool: diff --git a/tests/e2e/batches/test_bedrock_blank_s3_env_e2e.py b/tests/e2e/batches/test_bedrock_blank_s3_env_e2e.py deleted file mode 100644 index 77eb8427e59..00000000000 --- a/tests/e2e/batches/test_bedrock_blank_s3_env_e2e.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Live e2e pin for Bedrock batch create with blank AWS_S3_* env vars. - -Owns its own file (not test_batches_e2e.py) so the PR changed-file e2e gate -stays a single tiny file: this class boots its own gateway with -AWS_S3_ENCRYPTION_KEY_ID and AWS_S3_BUCKET_OWNER exported empty, then runs the -unified target_model_names upload + batch create lifecycle against real Bedrock. -""" - -from __future__ import annotations - -import json -from typing import Final - -import pytest -from batch_cleanup import cleanup_batch, cleanup_file -from batch_client import BatchClient, BatchCreateBody, BatchObject, FileObject -from bedrock_env_gateway import BedrockEnvGateway -from capabilities import is_managed_id -from e2e_http import FileUploadForm, require_successful_call, unwrap -from lifecycle import ResourceManager -from models import KeyGenerateBody - -pytestmark = pytest.mark.e2e - -CREATED_BATCH_STATUSES = {"validating", "in_progress", "finalizing"} -BLANK_S3_RAW_MODEL: Final = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" - - -def render_jsonl(model: str) -> bytes: - line = { - "custom_id": "req-1", - "method": "POST", - "url": "/v1/chat/completions", - "body": { - "model": model, - "messages": [{"role": "user", "content": "ping"}], - "max_tokens": 8, - }, - } - return (json.dumps(line) + "\n").encode() - - -def assert_file_object(file: FileObject, *, provider: str) -> None: - assert file.object == "file", f"file.object={file.object!r}" - assert file.purpose == "batch", f"file.purpose={file.purpose!r}" - assert file.bytes is not None, f"file.bytes={file.bytes!r}" - if provider != "bedrock": - assert file.bytes > 0, f"file.bytes={file.bytes!r}" - assert file.status, "file.status missing" - assert file.created_at is not None and file.created_at > 0, "file.created_at missing" - - -def assert_batch_object(batch: BatchObject) -> None: - assert batch.object == "batch", f"batch.object={batch.object!r}" - if batch.endpoint: - assert batch.endpoint == "/v1/chat/completions", f"batch.endpoint={batch.endpoint!r}" - assert batch.completion_window == "24h", f"window={batch.completion_window!r}" - assert batch.input_file_id, "batch.input_file_id missing" - assert batch.created_at is not None and batch.created_at > 0, "batch.created_at missing" - - -class TestBedrockBatchBlankS3EnvVars: - """Bedrock batch create with AWS_S3_* env vars exported but blank. - - Regression: a blank AWS_S3_ENCRYPTION_KEY_ID or AWS_S3_BUCKET_OWNER env var - resolved to "" and was serialized into the create-job request, which Bedrock - rejects. The owned gateway exports both vars empty, so the unified lifecycle - only passes when blank is treated as unset. - """ - - @pytest.mark.covers( - "llm.batches.bedrock.blank_s3_env.nonstream.works", - "llm.files.bedrock.upload.nonstream.works", - exercised_on=["batches", "files"], - ) - def test_unified_batch_create_ignores_blank_s3_env_vars(self, resources: ResourceManager) -> None: - gateway: Final = BedrockEnvGateway.start() - resources.defer(gateway.stop) - client: Final = BatchClient(proxy=gateway.proxy) - - key: Final = client.proxy.generate_key(KeyGenerateBody(models=[], user_id="e2e-test-user")) - resources.defer(lambda: client.proxy.delete_key(key)) - - file: Final = unwrap( - client.upload_file( - content=render_jsonl(BLANK_S3_RAW_MODEL), - form=FileUploadForm(purpose="batch", target_model_names="bedrock-blank-s3-batch"), - key=key, - ) - ) - resources.defer(lambda: cleanup_file(client, file.id, key=key)) - assert_file_object(file, provider="bedrock") - - created: Final = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) - assert created.status_code < 400, ( - f"blank AWS_S3_ENCRYPTION_KEY_ID / AWS_S3_BUCKET_OWNER must be treated as " - f"unset; Bedrock rejected the job: {created.body[:400]}" - ) - 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"blank-S3-env create via target_model_names must return a managed batch id, got {batch.id!r}" - ) - assert batch.status in CREATED_BATCH_STATUSES, ( - f"blank-S3-env batch has non-transitional status {batch.status!r}" - ) - assert_batch_object(batch) diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 7d334ed41ff..d09199ed138 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -25,7 +25,6 @@ - {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.blank_s3_env.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: blank_s3_env, streaming: nonstream, assertions: [works], source: "test_bedrock_blank_s3_env_e2e.py", rationale: "Bedrock batch create treats blank AWS_S3_ENCRYPTION_KEY_ID / AWS_S3_BUCKET_OWNER env vars as unset instead of serializing empty strings"} - {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"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index fec1934059c..8417b51360e 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -65,7 +65,6 @@ LlmCapability = Literal[ "assume_role", "basic", "batch_deployment", - "blank_s3_env", "code_interpreter", "count_tokens", "govcloud_partition", diff --git a/tests/integration/providers/test_bedrock_batch_blank_s3_env_wire.py b/tests/integration/providers/test_bedrock_batch_blank_s3_env_wire.py new file mode 100644 index 00000000000..9f5b03e6c19 --- /dev/null +++ b/tests/integration/providers/test_bedrock_batch_blank_s3_env_wire.py @@ -0,0 +1,222 @@ +import contextlib +import datetime +import json +import socket +import socketserver +import ssl +import threading +import uuid +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import pytest +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.x509.oid import NameOID +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel + +MODEL_ID: Final = "us.anthropic.claude-haiku-4-5-20251001-v1:0" +REGION: Final = "us-east-1" +BEDROCK_AUTHORITY: Final = f"bedrock.{REGION}.amazonaws.com:443" +BUCKET: Final = "integration-blank-s3-bucket" +ROLE_ARN: Final = "arn:aws:iam::123456789012:role/integration-batch-role" +JOB_ARN_PREFIX: Final = f"arn:aws:bedrock:{REGION}:123456789012:model-invocation-job/" +KMS_KEY: Final = f"arn:aws:kms:{REGION}:123456789012:key/integration-batch-key" +BUCKET_OWNER: Final = "123456789012" +SSE_HEADER_PREFIX: Final = "x-amz-server-side-encryption" + + +@dataclass(frozen=True, slots=True) +class ConnectProxy: + url: str + authorities: SimpleQueue[str] + + +class _DataConfig(BaseModel): + s3InputDataConfig: dict[str, str] + + +class _OutputConfig(BaseModel): + s3OutputDataConfig: dict[str, str] + + +class _CreateJob(BaseModel): + modelId: str + roleArn: str + inputDataConfig: _DataConfig + outputDataConfig: _OutputConfig + + +def _tls_context(directory: Path) -> ssl.SSLContext: + key: Final = ec.generate_private_key(ec.SECP256R1()) + now: Final = datetime.datetime.now(datetime.timezone.utc) + name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, BEDROCK_AUTHORITY.split(":")[0])]) + certificate: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(days=1)) + .not_valid_after(now + datetime.timedelta(days=1)) + .sign(key, hashes.SHA256()) + ) + certificate_file: Final = directory / "bedrock.pem" + key_file: Final = directory / "bedrock.key" + certificate_file.write_bytes(certificate.public_bytes(serialization.Encoding.PEM)) + key_file.write_bytes( + key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + ) + context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(certificate_file, key_file) + return context + + +def _pipe(source: socket.socket, sink: socket.socket) -> None: + with contextlib.suppress(OSError): + for chunk in iter(lambda: source.recv(65536), b""): + sink.sendall(chunk) + with contextlib.suppress(OSError): + sink.shutdown(socket.SHUT_WR) + + +@contextmanager +def bedrock_tunnel(destination: Wire) -> Generator[ConnectProxy, None, None]: + authorities: Final[SimpleQueue[str]] = SimpleQueue() + destination_port: Final = int(destination.url.rsplit(":", 1)[1]) + + class Tunnel(socketserver.StreamRequestHandler): + rbufsize = 0 + request: socket.socket + + def handle(self) -> None: + authority: Final = self.rfile.readline().decode().split()[1] + while self.rfile.readline() not in (b"\r\n", b""): + pass + authorities.put(authority) + if authority != BEDROCK_AUTHORITY: + self.wfile.write(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\n\r\n") + return + self.wfile.write(b"HTTP/1.1 200 Connection established\r\n\r\n") + self.request.settimeout(10) + with socket.create_connection(("127.0.0.1", destination_port), timeout=10) as upstream: + outbound: Final = threading.Thread(target=_pipe, args=(self.request, upstream)) + outbound.start() + _pipe(upstream, self.request) + outbound.join(timeout=12) + + with socketserver.ThreadingTCPServer(("127.0.0.1", 0), Tunnel) as server: + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) + thread.start() + try: + yield ConnectProxy(f"http://127.0.0.1:{server.server_address[1]}", authorities) + finally: + server.shutdown() + thread.join(timeout=6) + + +def s3_peer(request: Request) -> Reply: + assert request.method == "PUT" and request.target.startswith(f"/{BUCKET}/"), request.target + return Reply(body=b"") + + +def bedrock_peer(request: Request) -> Reply: + if request.method == "POST" and request.target == "/model-invocation-job": + return Reply(body=json.dumps({"jobArn": JOB_ARN_PREFIX + uuid.uuid4().hex}).encode()) + return Reply(status=404, body=b'{"message": "not scripted"}') + + +def _without_uri(config: Mapping[str, str]) -> dict[str, str]: + return {name: value for name, value in config.items() if name != "s3Uri"} + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize( + ("kms_key", "bucket_owner", "sse_headers", "input_fields", "output_fields"), + [ + pytest.param("", "", {}, {}, {}, id="blank"), + pytest.param( + KMS_KEY, + BUCKET_OWNER, + {SSE_HEADER_PREFIX: "aws:kms", f"{SSE_HEADER_PREFIX}-aws-kms-key-id": KMS_KEY}, + {"s3BucketOwner": BUCKET_OWNER}, + {"s3BucketOwner": BUCKET_OWNER, "s3EncryptionKeyId": KMS_KEY}, + id="set", + ), + ], +) +def test_unified_bedrock_batch_sends_s3_env_settings_only_when_they_are_non_blank( + gateway: Gateway, + tmp_path: Path, + kms_key: str, + bucket_owner: str, + sse_headers: Mapping[str, str], + input_fields: Mapping[str, str], + output_fields: Mapping[str, str], +) -> None: + environment: Final = { + "AWS_S3_ENCRYPTION_KEY_ID": kms_key, + "AWS_S3_BUCKET_OWNER": bucket_owner, + "SSL_VERIFY": "False", + "AWS_EC2_METADATA_DISABLED": "true", + } + with ( + wire_server(s3_peer) as s3, + wire_server(bedrock_peer, tls=_tls_context(tmp_path)) as bedrock, + bedrock_tunnel(bedrock) as tunnel, + owned_proxy(gateway, tmp_path, {**environment, "HTTPS_PROXY": tunnel.url}) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"bedrock/{MODEL_ID}", + api_key=None, + api_base=None, + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + aws_region_name=REGION, + s3_bucket_name=BUCKET, + s3_endpoint_url=s3.url, + aws_batch_role_arn=ROLE_ARN, + ) + line: Final = { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": model, "messages": [{"role": "user", "content": "ping"}], "max_tokens": 8}, + } + uploaded: Final = candidate.request_multipart( + "/v1/files", + {"purpose": "batch", "target_model_names": model}, + {"file": ("in.jsonl", (json.dumps(line) + "\n").encode(), "application/jsonl")}, + ) + assert uploaded.status_code == 200, uploaded.text + created: Final = candidate.request( + "POST", + "/v1/batches", + {"input_file_id": uploaded.json()["id"], "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + ) + assert created.status_code == 200, created.text + assert created.json()["object"] == "batch" and created.json()["status"] == "validating", created.text + + puts: Final = s3.drain() + assert len(puts) == 1, [put.target for put in puts] + assert { + name: value for name, value in puts[0].headers.items() if name.startswith(SSE_HEADER_PREFIX) + } == sse_headers + + assert BEDROCK_AUTHORITY in {tunnel.authorities.get_nowait() for _ in range(tunnel.authorities.qsize())} + jobs: Final = tuple(request for request in bedrock.drain() if request.method == "POST") + assert len(jobs) == 1, [job.target for job in jobs] + job: Final = _CreateJob.model_validate_json(jobs[0].body) + assert job.modelId == MODEL_ID and job.roleArn == ROLE_ARN + assert job.inputDataConfig.s3InputDataConfig["s3Uri"] == f"s3:/{puts[0].target}" + assert _without_uri(job.inputDataConfig.s3InputDataConfig) == input_fields + assert _without_uri(job.outputDataConfig.s3OutputDataConfig) == output_fields