mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
test(e2e): assert only litellm-owned batch behavior and move the blank S3 env pin to an integration test (#43321)
This commit is contained in:
yuneng-jiang
GitHub
parent
d18fcb09d6
commit
e53e67ede5
9 changed files with 319 additions and 282 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 "<no log written>"
|
||||
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)
|
||||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -65,7 +65,6 @@ LlmCapability = Literal[
|
|||
"assume_role",
|
||||
"basic",
|
||||
"batch_deployment",
|
||||
"blank_s3_env",
|
||||
"code_interpreter",
|
||||
"count_tokens",
|
||||
"govcloud_partition",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue