scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 145

fix(s3_v2): bound concurrent S3 uploads per flush and add opt-in JSONL batch files (#41258)

Browse source 
* fix(s3_v2): bound concurrent S3 uploads per flush and add opt-in JSONL batch files

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): keep failed uploads queued, parse env-backed flags, add integration coverage

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(s3_v2): type test helpers and honor constructor bound when config value is null

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(s3): annotate required casts for the type-discipline gate

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): keep tenant prefixes, stable retries and cold storage safety in batch file mode

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): cover root-level batch file keys for codecov patch target

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): audit matrix across chat, messages and responses surfaces with sink faults

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): read sink objects under the lock in the audit cells

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(s3_v2): rebind the retry queue instead of slicing in place

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): ignore stray non-POST requests in the surface upstream

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(s3_v2): keep fake upload state on the fake client instead of nonlocal counters

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: yucheng <yucheng@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-09-24 07:32:00 +00:00 • committed by GitHub
parent 3fb6f8740b
commit 09ebb28473
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 1665 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/constants.py
View file

@ -48,6 +48,7 @@ DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5))
DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10))
DEFAULT_S3_BATCH_SIZE: Final = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512))
DEFAULT_S3_MAX_CONCURRENT_UPLOADS: Final = int(os.getenv("DEFAULT_S3_MAX_CONCURRENT_UPLOADS", "16"))
# https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
MAX_S3_OBJECT_KEY_BYTES: Final = 1024
S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64

35
litellm/integrations/s3.py
View file

@ -20,7 +20,8 @@ from litellm.constants import (
)
from litellm.types.utils import StandardLoggingPayload
_S3_LOG_PROMPTS_ONLY: Final = TypeAdapter(bool)
_S3_BOOL: Final = TypeAdapter(bool)
_UPLOAD_BOUND: Final = TypeAdapter(int)
def resolve_s3_log_prompts_only(configured: object, environ: Mapping[str, str] | None = None) -> bool:
@ -29,12 +30,42 @@ def resolve_s3_log_prompts_only(configured: object, environ: Mapping[str, str] |
if raw is None or raw == "":
return False
try:
return _S3_LOG_PROMPTS_ONLY.validate_python(raw.strip() if isinstance(raw, str) else raw)
return _S3_BOOL.validate_python(raw.strip() if isinstance(raw, str) else raw)
except ValidationError:
verbose_logger.warning("s3 logging: s3_log_prompts_only=%r is not a boolean, logging prompts only", raw)
return True
def resolve_s3_max_concurrent_uploads(configured: object, fallback: int) -> int:
if configured is None or configured == "":
return fallback
try:
bound: Final = _UPLOAD_BOUND.validate_python(configured.strip() if isinstance(configured, str) else configured)
except ValidationError:
verbose_logger.warning(
"s3 logging: s3_max_concurrent_uploads=%r is not an integer, using %s", configured, fallback
)
return fallback
if bound < 1:
verbose_logger.warning(
"s3 logging: s3_max_concurrent_uploads=%r must be at least 1, using %s", configured, fallback
)
return fallback
return bound
def resolve_s3_batch_file_upload(configured: object) -> bool:
if configured is None or configured == "":
return False
try:
return _S3_BOOL.validate_python(configured.strip() if isinstance(configured, str) else configured)
except ValidationError:
verbose_logger.warning(
"s3 logging: s3_batch_file_upload=%r is not a boolean, keeping per-request objects", configured
)
return False
def prompts_only_payload(payload: StandardLoggingPayload) -> StandardLoggingPayload:
return {**payload, "response": None}

126
litellm/integrations/s3_v2.py
View file

@ -3,26 +3,33 @@ s3 Bucket Logging Integration
async_log_success_event: Processes the event, stores it in memory for DEFAULT_S3_FLUSH_INTERVAL_SECONDS seconds or until DEFAULT_S3_BATCH_SIZE and then flushes to s3
async_log_failure_event: Processes the event, stores it in memory for DEFAULT_S3_FLUSH_INTERVAL_SECONDS seconds or until DEFAULT_S3_BATCH_SIZE and then flushes to s3
NOTE 1: S3 does not provide a BATCH PUT API endpoint, so we create tasks to upload each element individually
NOTE 1: S3 does not provide a BATCH PUT API endpoint; by default each element is uploaded concurrently (bounded by s3_max_concurrent_uploads), or with s3_batch_file_upload the whole flush is written as one .jsonl file
"""
import asyncio
import time
from collections.abc import Mapping
from datetime import datetime
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final, cast
from urllib.parse import quote
from uuid import uuid4
import httpx
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS
from litellm.constants import (
DEFAULT_S3_BATCH_SIZE,
DEFAULT_S3_FLUSH_INTERVAL_SECONDS,
DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
)
from litellm.integrations.s3 import (
get_s3_object_download_filename,
get_s3_object_key,
prompts_only_payload,
resolve_s3_batch_file_upload,
resolve_s3_log_prompts_only,
resolve_s3_max_concurrent_uploads,
resolve_sse_params,
)
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
@ -43,7 +50,20 @@ if TYPE_CHECKING:
from botocore.credentials import Credentials
def _s3_key_parent(s3_object_key: str) -> str:
return s3_object_key.rsplit("/", 1)[0] if "/" in s3_object_key else ""
class S3BatchUploadError(Exception):
def __init__(self, failed: int, total: int) -> None:
self.failed = failed
self.total = total
super().__init__(f"{failed} of {total} S3 uploads failed; events kept in queue for the next flush")
class S3Logger(CustomBatchLogger, BaseAWSLLM):
preserve_events_added_during_flush = True
def __init__(
self,
s3_bucket_name: str | None = None,
@ -71,6 +91,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_server_side_encryption: str | None = None,
s3_sse_kms_key_id: str | None = None,
s3_log_prompts_only: bool | None = None,
s3_max_concurrent_uploads: int = DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
s3_batch_file_upload: bool = False,
s3_callback_params_override: dict | None = None,
**kwargs,
):
@ -112,7 +134,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_server_side_encryption=s3_server_side_encryption,
s3_sse_kms_key_id=s3_sse_kms_key_id,
s3_log_prompts_only=s3_log_prompts_only,
s3_max_concurrent_uploads=s3_max_concurrent_uploads,
s3_batch_file_upload=s3_batch_file_upload,
)
self._upload_semaphore = asyncio.Semaphore(self.s3_max_concurrent_uploads)
verbose_logger.debug("s3 logger using endpoint url %s", s3_endpoint_url)
# IMPORTANT
@ -168,6 +193,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_server_side_encryption: str | None = None,
s3_sse_kms_key_id: str | None = None,
s3_log_prompts_only: bool | None = None,
s3_max_concurrent_uploads: int = DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
s3_batch_file_upload: bool = False,
params_source: dict | None = None,
):
"""
@ -226,6 +253,16 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
params.get("s3_sse_kms_key_id") or s3_sse_kms_key_id,
)
configured_bound: Final = params.get("s3_max_concurrent_uploads")
self.s3_max_concurrent_uploads = resolve_s3_max_concurrent_uploads(
s3_max_concurrent_uploads if configured_bound is None or configured_bound == "" else configured_bound,
DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
)
self.s3_batch_file_upload = s3_batch_file_upload or resolve_s3_batch_file_upload(
params.get("s3_batch_file_upload")
)
def _build_object_url(self, s3_object_key: str) -> str:
"""
Build the exact URL that is both signed and sent, with the key percent-encoded once.
@ -347,7 +384,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
verbose_logger.exception("s3 Layer Error - %s", e)
self.handle_callback_failure(callback_name="S3Logger")
async def async_upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement):
async def async_upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement) -> bool:
try:
import base64
import hashlib
@ -364,7 +401,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
url: Final = self._build_object_url(batch_logging_element.s3_object_key)
# Convert JSON to string
json_string: Final = safe_dumps(batch_logging_element.payload)
json_string: Final = (
batch_logging_element.body
if batch_logging_element.body is not None
else safe_dumps(batch_logging_element.payload)
)
# Calculate SHA256 hash of the content
content_hash: Final = hashlib.sha256(json_string.encode("utf-8")).hexdigest()
@ -374,7 +415,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
# Prepare the request
headers: Final = {
"Content-Type": "application/json",
"Content-Type": batch_logging_element.content_type,
"Content-MD5": content_md5,
"x-amz-content-sha256": content_hash,
"Content-Language": "en",
@ -421,27 +462,72 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
except Exception as e:
verbose_logger.exception("Error uploading to s3: %s", e)
self.handle_callback_failure(callback_name="S3Logger")
return False
return True
async def async_send_batch(self):
async def async_send_batch(self) -> None:
"""
Sends runs from self.log_queue.
Sends runs from self.log_queue
Returns: None
Raises: Does not raise an exception, will only verbose_logger.exception()
Raises S3BatchUploadError when any upload failed; CustomBatchLogger.flush_queue
keeps the surviving queue entries for the next flush.
"""
verbose_logger.debug("s3_v2 logger - sending batch of %s", len(self.log_queue))
if not self.log_queue:
batch: Final = tuple(self.log_queue)
if not batch:
return
verbose_logger.debug("s3_v2 logger - sending batch of %s", len(batch))
#########################################################
# Flush the log queue to s3
# the log queue can be bounded by DEFAULT_S3_BATCH_SIZE
# see custom_batch_logger.py which triggers the flush
#########################################################
for payload in self.log_queue:
asyncio.create_task(self.async_upload_data_to_s3(payload))
uploads: Final = self._batch_file_elements(batch) if self._batch_file_mode_active() else batch
results: Final = await asyncio.gather(*(self._upload_bounded(element) for element in uploads))
failed: Final = tuple(element for element, ok in zip(uploads, results, strict=True) if not ok)
if not failed:
return
self.log_queue = [*failed, *self.log_queue[len(batch) :]]
raise S3BatchUploadError(failed=len(failed), total=len(uploads))
def _batch_file_mode_active(self) -> bool:
if not self.s3_batch_file_upload:
return False
if litellm.cold_storage_custom_logger == "s3_v2":
verbose_logger.warning(
"s3 logging: s3_batch_file_upload is ignored because s3_v2 is the cold storage logger; "
"per-request objects are required for spend log lookups"
)
return False
return True
async def _upload_bounded(self, element: s3BatchLoggingElement) -> bool:
async with self._upload_semaphore:
return await self.async_upload_data_to_s3(element)
def _batch_file_elements(self, batch: tuple[s3BatchLoggingElement, ...]) -> tuple[s3BatchLoggingElement, ...]:
now: Final = datetime.now(timezone.utc)
groups: Final = {
parent: tuple(
element for element in batch if element.body is None and _s3_key_parent(element.s3_object_key) == parent
)
for parent in sorted({_s3_key_parent(element.s3_object_key) for element in batch if element.body is None})
}
return tuple(element for element in batch if element.body is not None) + tuple(
self._build_batch_file_element(elements, parent, now) for parent, elements in groups.items()
)
def _build_batch_file_element(
self, elements: tuple[s3BatchLoggingElement, ...], parent: str, now: datetime
) -> s3BatchLoggingElement:
batch_name: Final = f"batch_{now.strftime('%H-%M-%S')}_{uuid4().hex}"
return s3BatchLoggingElement(
payload={},
body="\n".join(safe_dumps(element.payload) for element in elements),
content_type="application/x-ndjson",
s3_object_key=f"{parent}/{batch_name}.jsonl" if parent else f"{batch_name}.jsonl",
s3_object_download_filename=f"{batch_name}.jsonl",
)
def create_s3_batch_logging_element(
self,
@ -521,7 +607,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
url: Final = self._build_object_url(batch_logging_element.s3_object_key)
# Convert JSON to string
json_string: Final = safe_dumps(batch_logging_element.payload)
json_string: Final = (
batch_logging_element.body
if batch_logging_element.body is not None
else safe_dumps(batch_logging_element.payload)
)
# Calculate SHA256 hash of the content
content_hash: Final = hashlib.sha256(json_string.encode("utf-8")).hexdigest()
@ -531,7 +621,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
# Prepare the request
headers: Final = {
"Content-Type": "application/json",
"Content-Type": batch_logging_element.content_type,
"Content-MD5": content_md5,
"x-amz-content-sha256": content_hash,
"Content-Language": "en",

2
litellm/types/integrations/s3_v2.py
View file

@ -9,3 +9,5 @@ class s3BatchLoggingElement(BaseModel):
payload: dict
s3_object_key: str
s3_object_download_filename: str
body: str | None = None
content_type: str = "application/json"

344
tests/integration/observability/_s3_v2_support.py Normal file
View file

@ -0,0 +1,344 @@
import asyncio
import json
import threading
import time
from collections.abc import Mapping
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from types import MappingProxyType
from typing import Final
import anthropic
import openai
import yaml
from integration._support.client import Gateway, JsonValue, eventually, object_value
from integration._support.wire import Reply, Request
BUCKET: Final = "integration-bucket"
PREFIX: Final = "integration-logs"
@dataclass(slots=True)
class RecordingS3Sink:
"""Records every accepted PUT body by target, tracks peak concurrency, and can reject a leading
run of PUT attempts with a chosen status before accepting. Serves stored bodies back on GET."""
fail_attempts: int = 0
fail_until: float = 0.0
fail_status: int = 503
delay_seconds: float = 0.5
lock: threading.Lock = field(default_factory=threading.Lock)
in_flight: int = 0
peak: int = 0
attempts: int = 0
store: dict[str, bytes] = field(default_factory=dict) # mutable-ok: GET reads must see writes from earlier PUTs
def respond(self, request: Request) -> Reply:
if request.method == "GET":
body: Final = self.store.get(request.target)
if body is None:
return Reply(status=404)
return Reply(body=body)
assert request.method == "PUT", request.method
assert request.target.startswith(f"/{BUCKET}/{PREFIX}/"), request.target
with self.lock:
self.attempts += 1
if self.attempts <= self.fail_attempts or time.time() < self.fail_until:
return Reply(
status=self.fail_status,
body=b"<Error><Code>SinkFailure</Code></Error>",
content_type="application/xml",
)
self.in_flight += 1
self.peak = max(self.peak, self.in_flight)
self.store[request.target] = request.body
time.sleep(self.delay_seconds)
with self.lock:
self.in_flight -= 1
return Reply()
def objects(self) -> Mapping[str, bytes]:
with self.lock:
return MappingProxyType(dict(self.store))
def payloads(self) -> tuple[dict[str, JsonValue], ...]:
return tuple(object_value(json.loads(line)) for body in self.objects().values() for line in body.splitlines())
def s3_config(
path: Path, sink_url: str, extra: Mapping[str, JsonValue], settings: Mapping[str, JsonValue] | None = None
) -> Path:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"].update(
{
"callbacks": ["s3_v2"],
"s3_callback_params": {
"s3_bucket_name": BUCKET,
"s3_region_name": "us-east-1",
"s3_endpoint_url": sink_url,
"s3_path": PREFIX,
"s3_aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
"s3_aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
**extra,
},
**(settings or {}),
}
)
target: Final = path / "s3_v2.yaml"
target.write_text(yaml.safe_dump(config))
return target
def _chat_completion(identity: str) -> dict[str, JsonValue]:
return {
"id": identity,
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
}
def _chat_stream_frames(identity: str) -> tuple[bytes, ...]:
chunks: Final = (
{
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": None}],
},
{
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
},
)
return tuple(f"data: {json.dumps(chunk)}\n\n".encode() for chunk in chunks) + (b"data: [DONE]\n\n",)
def _messages_completion(identity: str) -> dict[str, JsonValue]:
return {
"id": identity,
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5-20250929",
"content": [{"type": "text", "text": "ok"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 11, "output_tokens": 4},
}
def _messages_stream_frames(identity: str) -> tuple[bytes, ...]:
events: Final = (
(
"message_start",
{
"type": "message_start",
"message": {
"id": identity,
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5-20250929",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 11, "output_tokens": 1},
},
},
),
(
"content_block_start",
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
),
(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}},
),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
(
"message_delta",
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}},
),
("message_stop", {"type": "message_stop"}),
)
return tuple(f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events)
def _responses_completion(identity: str) -> dict[str, JsonValue]:
return {
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": [
{
"type": "message",
"id": f"msg_{identity}",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
}
],
"usage": {"input_tokens": 11, "output_tokens": 4, "total_tokens": 15},
}
def _responses_stream_frames(identity: str) -> tuple[bytes, ...]:
events: Final = (
(
"response.created",
{
"type": "response.created",
"response": {**_responses_completion(identity), "status": "in_progress", "output": []},
},
),
(
"response.output_text.delta",
{
"type": "response.output_text.delta",
"item_id": f"msg_{identity}",
"output_index": 0,
"content_index": 0,
"delta": "ok",
},
),
("response.completed", {"type": "response.completed", "response": _responses_completion(identity)}),
)
return tuple(f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events)
def surface_reply(request: Request) -> Reply:
"""Scripted upstream that echoes the caller's marker string back as the response id."""
if request.method != "POST" or not request.body:
return Reply(status=404)
body: Final = json.loads(request.body)
if request.target.endswith("/chat/completions"):
identity: Final = body["messages"][0]["content"]
if body.get("stream"):
return Reply(content_type="text/event-stream", chunks=_chat_stream_frames(identity))
return Reply(body=json.dumps(_chat_completion(identity)).encode())
if request.target.endswith("/messages"):
identity_messages: Final = body["messages"][0]["content"]
if body.get("stream"):
return Reply(content_type="text/event-stream", chunks=_messages_stream_frames(identity_messages))
return Reply(body=json.dumps(_messages_completion(identity_messages)).encode())
assert request.target.endswith("/responses"), request.target
identity_responses: Final = body["input"]
if body.get("stream"):
return Reply(content_type="text/event-stream", chunks=_responses_stream_frames(identity_responses))
return Reply(body=json.dumps(_responses_completion(identity_responses)).encode())
SURFACES: Final = ("chat", "chat_stream", "messages", "messages_stream", "responses", "responses_stream")
def call_surface(
candidate: Gateway, surface: str, openai_model: str, anthropic_model: str, key: str, marker: str
) -> tuple[str, str | None]:
"""Drive one request through the given surface; return (client-visible response id, x-litellm-call-id)."""
base: Final = str(candidate.client.base_url).rstrip("/")
headers: Final = {"Authorization": f"Bearer {key}"}
if surface == "chat":
reply: Final = openai.OpenAI(base_url=f"{base}/v1", api_key=key).chat.completions.create(
model=openai_model,
messages=[{"role": "user", "content": marker}],
extra_body={"cache": {"no-cache": True}},
)
return reply.id, None
async def chat_stream() -> str:
stream = await openai.AsyncOpenAI(base_url=f"{base}/v1", api_key=key).chat.completions.create(
model=openai_model,
messages=[{"role": "user", "content": marker}],
stream=True,
extra_body={"cache": {"no-cache": True}},
)
seen = ""
async for chunk in stream:
seen = chunk.id # rebind-ok: the stream yields one chunk at a time
return seen
if surface == "chat_stream":
return asyncio.run(chat_stream()), None
if surface in ("messages", "messages_stream"):
client: Final = anthropic.Anthropic(base_url=base, api_key="anthropic-placeholder", default_headers=headers)
if surface == "messages":
reply_messages: Final = client.messages.create(
model=anthropic_model, max_tokens=16, messages=[{"role": "user", "content": marker}]
)
return reply_messages.id, None
with client.messages.stream(
model=anthropic_model, max_tokens=16, messages=[{"role": "user", "content": marker}]
) as stream:
final: Final = stream.get_final_message()
return final.id, None
if surface == "responses":
response: Final = candidate.request(
"POST",
"/v1/responses",
{"model": openai_model, "input": marker, "cache": {"no-cache": True}},
key=key,
)
assert response.status_code == 200, response.text
return str(response.json()["id"]), response.headers.get("x-litellm-call-id")
assert surface == "responses_stream", surface
with candidate.client.stream(
"POST",
"/v1/responses",
json={"model": openai_model, "input": marker, "stream": True},
headers=headers,
) as response:
text: Final = response.read().decode()
assert response.status_code == 200, text
call_id: Final = response.headers.get("x-litellm-call-id")
assert marker in text, text
return marker, call_id
def collect_payloads(sink: RecordingS3Sink, count: int, seconds: float = 60) -> tuple[dict[str, JsonValue], ...]:
"""Wait until `count` stored payload lines exist, then return every stored payload object."""
def delivered() -> int:
return sum(len(body.splitlines()) for body in sink.objects().values())
eventually(delivered, lambda total: total >= count, seconds=seconds)
return sink.payloads()
def mixed_burst(
candidate: Gateway, openai_model: str, anthropic_model: str, key: str, marker: str, per_surface: int = 8
) -> tuple[tuple[str, str | None], ...]:
"""Fire `per_surface` requests on every surface; returns (response id, x-litellm-call-id) per request."""
jobs: Final = tuple(
(surface, f"{marker}-{surface}-{index}") for surface in SURFACES for index in range(per_surface)
)
def call(job: tuple[str, str]) -> tuple[str, str | None]:
surface, identity = job
return call_surface(candidate, surface, openai_model, anthropic_model, key, identity)
with ThreadPoolExecutor(max_workers=48) as pool:
return tuple(pool.map(call, jobs))
def matched_ids(
payloads: tuple[dict[str, JsonValue], ...], answered: tuple[tuple[str, str | None], ...]
) -> frozenset[str]:
"""Every payload must be accountable to an answered request by response id or litellm_call_id."""
response_ids: Final = frozenset(observed for observed, _ in answered)
call_ids: Final = frozenset(call_id for _, call_id in answered if call_id is not None)
landed: Final = []
for payload in payloads:
if payload["id"] in response_ids:
landed.append(payload["id"])
continue
assert payload["litellm_call_id"] in call_ids, f"unmatched payload {payload['id']!r}"
landed.append(str(payload["id"]))
return frozenset(landed)

97
tests/integration/observability/test_s3_v2_flush_surfaces.py Normal file
View file

@ -0,0 +1,97 @@
import re
import uuid
from pathlib import Path
from typing import Final
import pytest
from _s3_v2_support import (
BUCKET,
PREFIX,
RecordingS3Sink,
collect_payloads,
matched_ids,
mixed_burst,
s3_config,
surface_reply,
)
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import wire_server
PER_REQUEST_KEY: Final = re.compile(rf"^/{BUCKET}/{PREFIX}/\d{{4}}-\d{{2}}-\d{{2}}/.+\.json$")
BATCH_KEY: Final = re.compile(
rf"^/{BUCKET}/{PREFIX}/\d{{4}}-\d{{2}}-\d{{2}}/batch_\d{{2}}-\d{{2}}-\d{{2}}_[0-9a-f]{{32}}\.jsonl$"
)
@pytest.mark.covers("other.observability.s3_v2.mixed_surface_burst_bounds_puts_one_object_per_response_id")
def test_s3_v2_mixed_surface_burst_bounds_puts_one_object_per_response_id(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3mix" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink()
with wire_server(surface_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = s3_config(tmp_path, bucket.url, {})
with (
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate,
candidate.scenario() as scenario,
):
openai_model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
anthropic_model: Final = scenario.model(
model="anthropic/claude-sonnet-4-5-20250929", api_base=provider.url, api_key="synthetic-provider-key"
)
key: Final = scenario.key(models=[openai_model, anthropic_model])
answered: Final = mixed_burst(candidate, openai_model, anthropic_model, key, marker)
payloads: Final = collect_payloads(sink, len(answered))
targets: Final = tuple(sink.objects())
assert sum(1 for r in provider.drain() if r.method == "POST") == 48
assert sink.peak <= 16, f"peak concurrent PUTs {sink.peak} exceeded the default bound"
assert all(PER_REQUEST_KEY.match(target) for target in targets), list(targets)
assert len(targets) == 48
assert matched_ids(payloads, answered)
@pytest.mark.covers("other.observability.s3_v2.mixed_surface_batch_writes_ndjson_lines_per_response_id")
def test_s3_v2_mixed_surface_batch_writes_ndjson_lines_per_response_id(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3mixb" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink()
with wire_server(surface_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": True})
with (
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate,
candidate.scenario() as scenario,
):
openai_model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
anthropic_model: Final = scenario.model(
model="anthropic/claude-sonnet-4-5-20250929", api_base=provider.url, api_key="synthetic-provider-key"
)
key: Final = scenario.key(models=[openai_model, anthropic_model])
answered: Final = mixed_burst(candidate, openai_model, anthropic_model, key, marker)
payloads: Final = collect_payloads(sink, len(answered))
targets: Final = tuple(sink.objects())
puts: Final = bucket.drain()
assert sum(1 for r in provider.drain() if r.method == "POST") == 48
assert all(BATCH_KEY.match(target) for target in targets), list(targets)
assert all(put.headers["content-type"] == "application/x-ndjson" for put in puts), [put.headers for put in puts]
assert matched_ids(payloads, answered)
assert len(payloads) == 48
@pytest.mark.covers("other.observability.s3_v2.sink_outage_mid_mixed_burst_recovers_every_response_id")
def test_s3_v2_sink_outage_mid_mixed_burst_recovers_every_response_id(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3mixo" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink(fail_attempts=30, fail_status=503, delay_seconds=0.2)
with wire_server(surface_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = s3_config(tmp_path, bucket.url, {})
with (
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate,
candidate.scenario() as scenario,
):
openai_model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
anthropic_model: Final = scenario.model(
model="anthropic/claude-sonnet-4-5-20250929", api_base=provider.url, api_key="synthetic-provider-key"
)
key: Final = scenario.key(models=[openai_model, anthropic_model])
answered: Final = mixed_burst(candidate, openai_model, anthropic_model, key, marker)
payloads: Final = collect_payloads(sink, len(answered), seconds=90)
assert sum(1 for r in provider.drain() if r.method == "POST") == 48
assert matched_ids(payloads, answered)
assert len(payloads) == 48, "a stored id was overwritten or duplicated"

630
tests/integration/observability/test_s3_v2_upload_fanout.py Normal file
View file

@ -0,0 +1,630 @@
import json
import re
import threading
import time
import uuid
from collections.abc import Mapping
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from _s3_v2_support import RecordingS3Sink, collect_payloads
from _s3_v2_support import s3_config as _recording_s3_config
from integration._support.client import Gateway, JsonValue, eventually
from integration._support.process import group_members, owned_proxy, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
BUCKET: Final = "integration-bucket"
PREFIX: Final = "integration-logs"
REQUESTS: Final = 64
PUT_DELAY_SECONDS: Final = 0.5
@dataclass(slots=True)
class S3Sink:
"""Accepts every PUT after a fixed delay and records the peak number of PUTs in flight."""
lock: threading.Lock = field(default_factory=threading.Lock)
in_flight: int = 0
peak: int = 0
def respond(self, request: Request) -> Reply:
assert request.method == "PUT", request.method
assert request.target.startswith(f"/{BUCKET}/{PREFIX}/"), request.target
with self.lock:
self.in_flight += 1
self.peak = max(self.peak, self.in_flight)
time.sleep(PUT_DELAY_SECONDS)
with self.lock:
self.in_flight -= 1
return Reply()
def _chat_reply(request: Request) -> Reply:
if request.method != "POST" or not request.body:
return Reply(status=404)
text: Final = json.loads(request.body)["messages"][0]["content"]
return Reply(
body=json.dumps(
{
"id": text,
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
}
).encode()
)
def _s3_config(path: Path, sink_url: str, extra: Mapping[str, JsonValue]) -> Path:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["litellm_settings"].update(
{
"callbacks": ["s3_v2"],
"s3_callback_params": {
"s3_bucket_name": BUCKET,
"s3_region_name": "us-east-1",
"s3_endpoint_url": sink_url,
"s3_path": PREFIX,
"s3_aws_access_key_id": "AKIAIOSFODNN7EXAMPLE",
"s3_aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
**extra,
},
}
)
target: Final = path / "s3_v2.yaml"
target.write_text(yaml.safe_dump(config))
return target
def _burst(candidate: Gateway, model: str, key: str, marker: str) -> frozenset[str]:
ids: Final = tuple(f"{marker}-{index}" for index in range(REQUESTS))
def request(identity: str) -> str:
response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": identity}], "cache": {"no-cache": True}},
key=key,
)
assert response.status_code == 200, response.text
return response.json()["id"]
with ThreadPoolExecutor(max_workers=32) as pool:
returned: Final = frozenset(pool.map(request, ids))
assert returned == frozenset(ids)
return returned
def _collect(bucket: Wire, count_lines: bool, expected: int) -> tuple[Request, ...]:
puts: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier PUTs
def delivered() -> int:
puts.extend(bucket.drain())
return sum(len(put.body.splitlines()) if count_lines else 1 for put in puts)
eventually(delivered, lambda total: total >= expected, seconds=30)
return tuple(puts)
PER_REQUEST_KEY: Final = re.compile(rf"^/{BUCKET}/{PREFIX}/\d{{4}}-\d{{2}}-\d{{2}}/.+\.json$")
BATCH_KEY: Final = re.compile(
rf"^/{BUCKET}/{PREFIX}/\d{{4}}-\d{{2}}-\d{{2}}/batch_\d{{2}}-\d{{2}}-\d{{2}}_[0-9a-f]{{32}}\.jsonl$"
)
@pytest.mark.covers("other.observability.s3_v2.flush_bounds_concurrent_puts_to_default_and_keeps_every_log")
def test_s3_v2_flush_bounds_concurrent_puts_to_the_default_of_sixteen(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3fan" + uuid.uuid4().hex[:8]
sink: Final = S3Sink()
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {})
with (
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
ids: Final = _burst(candidate, model, key, marker)
puts: Final = _collect(bucket, count_lines=False, expected=REQUESTS)
assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS
assert sink.peak <= 16, f"peak concurrent PUTs {sink.peak} exceeded the default bound for {REQUESTS} queued logs"
assert all(PER_REQUEST_KEY.match(put.target) for put in puts), [put.target for put in puts]
assert frozenset(json.loads(put.body)["id"] for put in puts) == ids
assert len({put.target for put in puts}) == REQUESTS
@pytest.mark.covers("other.observability.s3_v2.configured_bound_and_env_backed_false_keeps_per_request_objects")
def test_s3_v2_honors_configured_bound_and_env_backed_false_batch_flag(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3cap" + uuid.uuid4().hex[:8]
sink: Final = S3Sink()
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(
tmp_path,
bucket.url,
{"s3_max_concurrent_uploads": 4, "s3_batch_file_upload": "os.environ/INTEGRATION_S3_BATCH_FILE_UPLOAD"},
)
with (
owned_proxy(
gateway,
tmp_path,
{"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3", "INTEGRATION_S3_BATCH_FILE_UPLOAD": "false"},
config=config,
) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
ids: Final = _burst(candidate, model, key, marker)
puts: Final = _collect(bucket, count_lines=False, expected=REQUESTS)
assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS
assert sink.peak <= 4, f"peak concurrent PUTs {sink.peak} exceeded s3_max_concurrent_uploads=4"
assert all(PER_REQUEST_KEY.match(put.target) for put in puts), [put.target for put in puts]
assert frozenset(json.loads(put.body)["id"] for put in puts) == ids
@pytest.mark.covers("other.observability.s3_v2.batch_file_upload_writes_one_ndjson_object_per_flush")
def test_s3_v2_batch_file_upload_writes_one_jsonl_object_per_flush(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3jsonl" + uuid.uuid4().hex[:8]
sink: Final = S3Sink()
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": True})
with (
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
ids: Final = _burst(candidate, model, key, marker)
puts: Final = _collect(bucket, count_lines=True, expected=REQUESTS)
assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS
assert len(puts) <= 2, f"{len(puts)} PUTs for {REQUESTS} logs; batch mode must write one object per flush"
assert all(BATCH_KEY.match(put.target) for put in puts), [put.target for put in puts]
assert all(put.headers["content-type"] == "application/x-ndjson" for put in puts), [put.headers for put in puts]
lines: Final = tuple(line for put in puts for line in put.body.decode().splitlines())
assert frozenset(json.loads(line)["id"] for line in lines) == ids
assert len(lines) == REQUESTS
@pytest.mark.covers("other.observability.s3_v2.batch_file_upload_keeps_team_prefix_in_object_key")
def test_s3_v2_batch_file_upload_keeps_team_alias_prefix(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3team" + uuid.uuid4().hex[:8]
team_alias: Final = f"alpha-{uuid.uuid4().hex[:8]}"
team_batch_key: Final = re.compile(
rf"^/{BUCKET}/{PREFIX}/{team_alias}/\d{{4}}-\d{{2}}-\d{{2}}/batch_\d{{2}}-\d{{2}}-\d{{2}}_[0-9a-f]{{32}}\.jsonl$"
)
sink: Final = S3Sink()
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": True, "s3_use_team_prefix": True})
with (
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
team: Final = scenario.team(team_alias=team_alias, models=[model])
key: Final = scenario.key(team_id=team, models=[model])
ids: Final = _burst(candidate, model, key, marker)
puts: Final = _collect(bucket, count_lines=True, expected=REQUESTS)
assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS
assert len(puts) >= 1
assert all(team_batch_key.match(put.target) for put in puts), [put.target for put in puts]
lines: Final = tuple(line for put in puts for line in put.body.decode().splitlines())
assert frozenset(json.loads(line)["id"] for line in lines) == ids
assert len(lines) == REQUESTS
@pytest.mark.covers("other.observability.s3_v2.upstream_failure_events_land_alongside_successes")
def test_s3_v2_upstream_failure_events_land_alongside_successes(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3fail" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink()
def provider(request: Request) -> Reply:
text: Final = json.loads(request.body)["messages"][0]["content"]
if text.endswith("-fail"):
return Reply(
status=401,
body=b'{"error": {"message": "synthetic upstream rejection", "code": "synthetic_401"}}',
)
return _chat_reply(request)
with wire_server(provider) as upstream, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {})
with (
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=upstream.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
success_ids: Final = tuple(f"{marker}-{index}" for index in range(8))
failure_ids: Final = tuple(f"{marker}-{index}-fail" for index in range(4))
def send(identity: str) -> httpx.Response:
return candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": identity}], "cache": {"no-cache": True}},
key=key,
)
with ThreadPoolExecutor(max_workers=12) as pool:
responses: Final = tuple(pool.map(send, (*success_ids, *failure_ids)))
ok: Final = responses[:8]
rejected: Final = responses[8:]
assert all(response.status_code == 200 for response in ok), [r.text for r in ok]
assert tuple(response.json()["id"] for response in ok) == success_ids
for response in rejected:
assert response.status_code in (400, 401), response.status_code
assert "synthetic upstream rejection" in response.text, response.text
failure_call_ids: Final = frozenset(response.headers["x-litellm-call-id"] for response in rejected)
payloads: Final = collect_payloads(sink, len(success_ids) + len(failure_ids))
assert len(upstream.drain()) == len(success_ids) + len(failure_ids)
delivered: Final = frozenset(payload["id"] for payload in payloads if payload["status"] == "success")
assert delivered == frozenset(success_ids)
failures: Final = tuple(payload for payload in payloads if payload["status"] == "failure")
assert len(failures) == len(failure_ids)
assert frozenset(payload["litellm_call_id"] for payload in failures) == failure_call_ids
assert all("synthetic upstream rejection" in json.dumps(payload["error_information"]) for payload in failures)
@pytest.mark.covers("other.observability.s3_v2.invalid_or_empty_bound_falls_back_to_sixteen")
@pytest.mark.parametrize(
("bad", "warns"),
[
pytest.param("abc", True, id="non_integer"),
pytest.param(0, True, id="below_one"),
pytest.param("", False, id="empty"),
],
)
def test_s3_v2_invalid_or_empty_bound_falls_back_to_sixteen(
gateway: Gateway, tmp_path: Path, bad: JsonValue, warns: bool
) -> None:
marker: Final = "s3bound" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink()
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {"s3_max_concurrent_uploads": bad})
with (
owned_proxy_process(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
ids: Final = _burst(owned.gateway, model, key, marker)
payloads: Final = collect_payloads(sink, REQUESTS)
if warns:
eventually(
lambda: owned.log.read_text(),
lambda text: "s3_max_concurrent_uploads" in text,
seconds=15,
)
else:
assert "s3_max_concurrent_uploads" not in owned.log.read_text()
assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS
assert sink.peak <= 16, f"peak concurrent PUTs {sink.peak} exceeded the fallback bound"
assert frozenset(payload["id"] for payload in payloads) == ids
@pytest.mark.covers("other.observability.s3_v2.sink_rejection_requeues_and_delivers_every_id_once")
def test_s3_v2_sink_rejection_requeues_and_delivers_every_id_once(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3deny" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink(fail_status=403, delay_seconds=0.2)
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {})
with (
owned_proxy_process(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
sink.fail_until = time.time() + 10
ids: Final = _burst(owned.gateway, model, key, marker)
payloads: Final = collect_payloads(sink, REQUESTS, seconds=90)
eventually(
lambda: owned.log.read_text(),
lambda text: "S3BatchUploadError" in text,
seconds=15,
)
readiness: Final = owned.gateway.client.get("/health/readiness")
assert readiness.status_code == 200, readiness.text
assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS
assert len(sink.objects()) == REQUESTS
assert frozenset(payload["id"] for payload in payloads) == ids
@pytest.mark.covers("other.observability.s3_v2.batch_retry_resends_identical_key_and_body")
def test_s3_v2_batch_retry_resends_identical_key_and_body(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3retry" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink(fail_status=500, delay_seconds=0.2)
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": True})
with (
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
sink.fail_until = time.time() + 8
ids: Final = _burst(candidate, model, key, marker)
payloads: Final = collect_payloads(sink, REQUESTS, seconds=90)
puts: Final = bucket.drain()
assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS
by_target: Final = {}
for put in puts:
by_target.setdefault(put.target, set()).add(put.body) # mutable-ok: grouping attempts seen so far per target
assert all(len(bodies) == 1 for bodies in by_target.values()), "a retried batch PUT changed key or body"
assert max(sum(1 for put in puts if put.target == target) for target in by_target) >= 2, "no retried PUT observed"
assert frozenset(payload["id"] for payload in payloads) == ids
assert len(payloads) == REQUESTS
@pytest.mark.covers("other.observability.s3_v2.unknown_model_rejection_keeps_other_requests_logging")
def test_s3_v2_unknown_model_rejection_keeps_other_requests_logging(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3ghost" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink()
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {})
with (
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
ghost: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": f"ghost-{uuid.uuid4().hex}", "messages": [{"role": "user", "content": "hi"}]},
key=key,
)
assert ghost.status_code in (400, 403, 404), ghost.text
ids: Final = _burst(candidate, model, key, marker)
eventually(
lambda: frozenset(payload["id"] for payload in sink.payloads()),
lambda landed: ids <= landed,
seconds=90,
)
payloads: Final = sink.payloads()
assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS
assert ids <= frozenset(payload["id"] for payload in payloads)
extras: Final = tuple(payload for payload in payloads if payload["id"] not in ids)
assert all(payload["status"] == "failure" for payload in extras), extras
@pytest.mark.covers("other.observability.s3_v2.batch_flag_ignored_when_s3_v2_is_cold_storage_logger")
def test_s3_v2_batch_flag_ignored_when_s3_v2_is_cold_storage_logger(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3cold" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink()
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _recording_s3_config(
tmp_path,
bucket.url,
{"s3_batch_file_upload": True},
{"cold_storage_custom_logger": "s3_v2"},
)
with (
owned_proxy_process(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
response: Final = owned.gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": marker}], "cache": {"no-cache": True}},
key=key,
)
assert response.status_code == 200, response.text
request_id: Final = str(response.json()["id"])
payloads: Final = collect_payloads(sink, 1)
assert all(PER_REQUEST_KEY.match(target) for target in sink.objects()), list(sink.objects())
eventually(
lambda: owned.log.read_text(),
lambda text: "s3_batch_file_upload is ignored because s3_v2 is the cold storage logger" in text,
seconds=15,
)
spend: Final = eventually(
lambda: owned.gateway.request("GET", f"/spend/logs/ui/{request_id}"),
lambda reply: reply.status_code == 200 and bool((reply.json() or {}).get("messages")),
seconds=60,
)
assert spend.status_code == 200, spend.text
body: Final = spend.json()
assert body["messages"], spend.text
assert body["response"], spend.text
assert payloads[0]["id"] == request_id
@pytest.mark.covers("other.observability.s3_v2.identical_requests_land_distinct_objects")
def test_s3_v2_identical_requests_land_distinct_objects(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3same" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink()
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {})
with (
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
def send(_: int) -> str:
response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": marker}], "cache": {"no-cache": True}},
key=key,
)
assert response.status_code == 200, response.text
return str(response.json()["id"])
with ThreadPoolExecutor(max_workers=16) as pool:
returned: Final = frozenset(pool.map(send, range(16)))
payloads: Final = collect_payloads(sink, 16)
assert sum(1 for r in provider.drain() if r.method == "POST") == 16
assert returned == {marker}, "the upstream echo keeps the same id for identical requests"
assert len(sink.objects()) == 16, "identical requests must still land as distinct objects"
assert all(payload["id"] == marker for payload in payloads)
@pytest.mark.covers("other.observability.s3_v2.two_workers_bound_and_deliver_every_id")
def test_s3_v2_two_workers_bound_and_deliver_every_id(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3work" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink()
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {})
with (
owned_proxy(
gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config, workers=2
) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
ids: Final = _burst(candidate, model, key, marker)
payloads: Final = collect_payloads(sink, REQUESTS)
assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS
assert sink.peak <= 32, f"peak concurrent PUTs {sink.peak} exceeded two workers at the default bound"
assert len(sink.objects()) == REQUESTS
assert frozenset(payload["id"] for payload in payloads) == ids
@pytest.mark.covers("other.observability.s3_v2.slow_sink_never_duplicates_or_stalls_readiness")
def test_s3_v2_slow_sink_never_duplicates_or_stalls_readiness(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3slow" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink(delay_seconds=1.5)
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": True})
with (
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "1"}, config=config) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
ids: Final = _burst(candidate, model, key, marker)
def delivered() -> int:
readiness: Final = candidate.client.get("/health/readiness")
assert readiness.status_code == 200, readiness.text
return sum(len(body.splitlines()) for body in sink.objects().values())
eventually(delivered, lambda total: total >= REQUESTS, seconds=90)
payloads: Final = sink.payloads()
puts: Final = bucket.drain()
targets: Final = tuple(put.target for put in puts)
assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS
assert len(set(targets)) == len(targets), "the same object was PUT more than once"
assert frozenset(payload["id"] for payload in payloads) == ids
assert len(payloads) == REQUESTS
@pytest.mark.covers("other.observability.s3_v2.worker_kill_mid_burst_keeps_surviving_deliveries")
def test_s3_v2_worker_kill_mid_burst_keeps_surviving_deliveries(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3kill" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink()
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {})
with (
owned_proxy_process(
gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config, workers=2
) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
sent: Final = tuple(f"{marker}-{index}" for index in range(REQUESTS))
def send(identity: str) -> tuple[str, bool]:
try:
response: Final = owned.gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": identity}],
"cache": {"no-cache": True},
},
key=key,
)
except Exception:
return identity, False
return identity, response.status_code == 200
with ThreadPoolExecutor(max_workers=32) as pool:
futures: Final = tuple(pool.submit(send, identity) for identity in sent)
time.sleep(0.5)
children: Final = tuple(
process for process in group_members(owned.process.pid) if process.pid != owned.process.pid
)
assert children, "no worker children found to kill"
children[0].kill()
results: Final = tuple(future.result() for future in futures)
survivors: Final = frozenset(identity for identity, ok in results if ok)
assert survivors, "no request survived the worker kill"
readiness: Final = owned.gateway.client.get("/health/readiness")
assert readiness.status_code == 200, readiness.text
payloads: Final = collect_payloads(sink, len(survivors), seconds=90)
landed: Final = frozenset(payload["id"] for payload in payloads)
assert survivors <= landed, "an id whose response succeeded never landed"
assert landed <= frozenset(sent), "an id that was never sent landed"
@pytest.mark.covers("other.observability.s3_v2.sigterm_mid_burst_loses_only_inflight_without_duplicates")
def test_s3_v2_sigterm_mid_burst_loses_only_inflight_without_duplicates(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3term" + uuid.uuid4().hex[:8]
sink: Final = RecordingS3Sink()
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {})
owned: Final = owned_proxy_process(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config)
candidate_owned: Final = owned.__enter__()
try:
created: Final = candidate_owned.gateway.post(
"/model/new",
{
"model_name": f"integration-{marker}",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "synthetic-provider-key",
"api_base": provider.url + "/v1",
},
"model_info": {},
},
)
model: Final = str(created["model_name"])
key: Final = str(candidate_owned.gateway.post("/key/generate", {"models": [model]})["key"])
sent: Final = tuple(f"{marker}-{index}" for index in range(REQUESTS))
def send(identity: str) -> tuple[str, bool]:
try:
response: Final = candidate_owned.gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": identity}],
"cache": {"no-cache": True},
},
key=key,
)
except Exception:
return identity, False
return identity, response.status_code == 200
with ThreadPoolExecutor(max_workers=32) as pool:
futures: Final = tuple(pool.submit(send, identity) for identity in sent)
time.sleep(0.5)
candidate_owned.process.terminate()
results: Final = tuple(future.result() for future in futures)
candidate_owned.process.wait(timeout=30)
finally:
owned.__exit__(None, None, None)
answered: Final = frozenset(identity for identity, ok in results if ok)
landed: Final = frozenset(payload["id"] for payload in sink.payloads())
assert landed <= answered, (
"a delivered object has no matching answered request; lost in-flight ids are expected, extras are not"
)
targets: Final = tuple(sink.objects())
assert len(set(targets)) == len(targets), "the same object was PUT more than once"

450
tests/test_litellm/integrations/test_s3_v2.py
View file

@ -2468,3 +2468,453 @@ def test_prompts_only_toggle_is_exposed_to_admin_ui_for_both_s3_callbacks(callba
from litellm.integrations.custom_logger import CustomLogger
assert "S3_LOG_PROMPTS_ONLY" in CustomLogger.get_callback_env_vars(callback_name)
def _element(payload: dict[str, object], key_suffix: str) -> s3BatchLoggingElement:
return s3BatchLoggingElement(
s3_object_key=f"2025-09-14/test-{key_suffix}.json",
payload=payload,
s3_object_download_filename=f"test-{key_suffix}.json",
)
def _ok_response() -> MagicMock:
response = MagicMock()
response.status_code = 200
response.raise_for_status = MagicMock()
return response
class _CountingPut:
def __init__(self) -> None:
self.in_flight = 0
self.peak = 0
self.calls = 0
async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
self.in_flight += 1
self.peak = max(self.peak, self.in_flight)
self.calls += 1
await asyncio.sleep(0.01)
self.in_flight -= 1
return _ok_response()
class _RecordingPut:
def __init__(self) -> None:
self.calls: tuple[tuple[str, str | None, dict[str, str] | None], ...] = ()
async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
self.calls = (*self.calls, (url, data, headers))
return _ok_response()
class _LateAppendingPut:
def __init__(self, logger: S3Logger, element: s3BatchLoggingElement, fail_first: bool = False) -> None:
self.logger = logger
self.element = element
self.fail_first = fail_first
self.appended = False
async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
if not self.appended:
self.appended = True
self.logger.log_queue.append(self.element)
if self.fail_first:
return _failure_response()
return _ok_response()
class _FailOnSuffixPut:
def __init__(self, suffixes: tuple[str, ...]) -> None:
self.failing = True
self.suffixes = suffixes
async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
if self.failing and url.endswith(self.suffixes):
return _failure_response()
return _ok_response()
class _FailUntilClearedPut:
def __init__(self) -> None:
self.failing = True
self.calls: tuple[tuple[str, str | None], ...] = ()
async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
self.calls = (*self.calls, (url, data))
if self.failing:
return _failure_response()
return _ok_response()
@pytest.mark.asyncio
async def test_async_send_batch_bounds_concurrent_uploads() -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_max_concurrent_uploads=4,
)
put = _CountingPut()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = put
logger.log_queue = [_element({"i": i}, f"{i}") for i in range(40)]
await logger.async_send_batch()
assert put.peak == 4
assert put.calls == 40
@pytest.mark.asyncio
async def test_async_send_batch_uploads_single_jsonl_file() -> None:
import json
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_batch_file_upload=True,
)
put = _RecordingPut()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = put
payloads = [{"id": "req-1"}, {"id": "req-2"}, {"id": "req-3"}]
logger.log_queue = [_element(payload, f"{i}") for i, payload in enumerate(payloads)]
await logger.async_send_batch()
assert len(put.calls) == 1
url, data, headers = put.calls[0]
assert url.endswith(".jsonl")
assert data is not None
assert headers is not None
assert [json.loads(line) for line in data.splitlines()] == payloads
assert headers["Content-Type"] == "application/x-ndjson"
@pytest.mark.asyncio
async def test_flush_queue_preserves_events_added_during_upload() -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
)
late_element = _element({"id": "late"}, "late")
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = _LateAppendingPut(logger, late_element)
logger.log_queue = [_element({"id": "first"}, "first")]
await logger.flush_queue()
assert logger.log_queue == [late_element]
def _override_logger(**overrides: object) -> S3Logger:
return S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_callback_params_override=overrides,
)
def test_env_backed_false_string_keeps_per_request_uploads() -> None:
assert _override_logger(s3_batch_file_upload="false").s3_batch_file_upload is False
assert _override_logger(s3_batch_file_upload="true").s3_batch_file_upload is True
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_batch_file_upload=True,
s3_callback_params_override={"s3_batch_file_upload": "false"},
)
assert logger.s3_batch_file_upload is True
@pytest.mark.parametrize("bad", [0, -3, "0", "abc", ""])
def test_invalid_concurrency_falls_back_to_default(bad: object) -> None:
from litellm.constants import DEFAULT_S3_MAX_CONCURRENT_UPLOADS
logger = _override_logger(s3_max_concurrent_uploads=bad)
assert logger.s3_max_concurrent_uploads == DEFAULT_S3_MAX_CONCURRENT_UPLOADS
assert logger._upload_semaphore._value == DEFAULT_S3_MAX_CONCURRENT_UPLOADS
def test_env_backed_concurrency_string_is_parsed() -> None:
logger = _override_logger(s3_max_concurrent_uploads="4")
assert logger.s3_max_concurrent_uploads == 4
assert logger._upload_semaphore._value == 4
@pytest.mark.parametrize("empty", [None, ""])
def test_empty_config_concurrency_falls_back_to_constructor_value(empty: object) -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_max_concurrent_uploads=4,
s3_callback_params_override={"s3_max_concurrent_uploads": empty},
)
assert logger.s3_max_concurrent_uploads == 4
assert logger._upload_semaphore._value == 4
def _failure_response() -> MagicMock:
response = MagicMock()
response.status_code = 400
response.raise_for_status = MagicMock(side_effect=Exception("s3 rejected the object"))
return response
@pytest.mark.asyncio
async def test_failed_uploads_stay_queued_for_next_flush() -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
)
elements = [_element({"i": i}, f"{i}") for i in range(5)]
put = _FailOnSuffixPut(("test-2.json", "test-4.json"))
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = put
logger.log_queue = list(elements)
await logger.flush_queue()
assert logger.log_queue == [elements[2], elements[4]]
put.failing = False
await logger.flush_queue()
assert logger.log_queue == []
@pytest.mark.asyncio
async def test_batch_file_upload_failure_keeps_whole_batch() -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_batch_file_upload=True,
)
put = _FailUntilClearedPut()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = put
elements = [_element({"i": i}, f"{i}") for i in range(3)]
logger.log_queue = list(elements)
await logger.flush_queue()
assert len(put.calls) == 1
assert len(logger.log_queue) == 1
assert logger.log_queue[0].body == "\n".join(json.dumps(element.payload) for element in elements)
@pytest.mark.asyncio
async def test_events_appended_during_failed_flush_survive() -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
)
late = _element({"id": "late"}, "late")
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = _LateAppendingPut(logger, late, fail_first=True)
first = _element({"id": "first"}, "first")
logger.log_queue = [first]
await logger.flush_queue()
assert logger.log_queue == [first, late]
@pytest.mark.asyncio
async def test_batch_file_key_shape() -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_path="logs",
s3_batch_file_upload=True,
)
put = _RecordingPut()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = put
logger.log_queue = [_element({"id": "req-1"}, "0")]
await logger.async_send_batch()
((url, _data, headers),) = put.calls
assert headers is not None
assert re.search(r".*/2025-09-14/batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl$", url)
assert headers["Content-Disposition"].endswith('.jsonl"')
@pytest.mark.asyncio
async def test_batch_file_groups_raw_elements_by_key_parent() -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_batch_file_upload=True,
)
put = _RecordingPut()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = put
alpha = s3BatchLoggingElement(
s3_object_key="logs/alpha/2026-01-01/a.json", payload={"id": "a"}, s3_object_download_filename="a.json"
)
beta = s3BatchLoggingElement(
s3_object_key="logs/beta/2026-01-01/b.json", payload={"id": "b"}, s3_object_download_filename="b.json"
)
plain = s3BatchLoggingElement(
s3_object_key="logs/2026-01-01/c.json", payload={"id": "c"}, s3_object_download_filename="c.json"
)
root = s3BatchLoggingElement(
s3_object_key="solo.json", payload={"id": "d"}, s3_object_download_filename="solo.json"
)
logger.log_queue = [alpha, beta, plain, root]
await logger.async_send_batch()
assert len(put.calls) == 4
by_parent = {
re.sub(r"(^|/)batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl$", "", url.split(".com/", 1)[-1]): (url, data)
for url, data, _headers in put.calls
}
assert sorted(by_parent) == ["", "logs/2026-01-01", "logs/alpha/2026-01-01", "logs/beta/2026-01-01"]
assert [line for line in by_parent[""][1].splitlines()] == [json.dumps({"id": "d"})]
assert [line for line in by_parent["logs/alpha/2026-01-01"][1].splitlines()] == [json.dumps({"id": "a"})]
assert [line for line in by_parent["logs/beta/2026-01-01"][1].splitlines()] == [json.dumps({"id": "b"})]
assert [line for line in by_parent["logs/2026-01-01"][1].splitlines()] == [json.dumps({"id": "c"})]
@pytest.mark.asyncio
async def test_failed_batch_file_is_requeued_and_resent_unchanged() -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_batch_file_upload=True,
)
put = _FailUntilClearedPut()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = put
logger.log_queue = [_element({"i": i}, f"{i}") for i in range(3)]
await logger.flush_queue()
assert len(logger.log_queue) == 1
assert logger.log_queue[0].body is not None
assert logger.log_queue[0].s3_object_key.endswith(".jsonl")
put.failing = False
await logger.flush_queue()
assert logger.log_queue == []
assert len(put.calls) == 2
assert put.calls[0] == put.calls[1]
@pytest.mark.asyncio
async def test_elements_appended_after_failed_batch_file_get_their_own_file() -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_batch_file_upload=True,
)
put = _FailUntilClearedPut()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = put
logger.log_queue = [_element({"id": "first"}, "first")]
await logger.flush_queue()
late = _element({"id": "late"}, "late")
logger.log_queue.append(late)
put.failing = False
await logger.flush_queue()
assert logger.log_queue == []
assert len(put.calls) == 3
assert put.calls[0] == put.calls[1]
assert put.calls[2][0] != put.calls[0][0]
assert put.calls[2][1] == json.dumps({"id": "late"})
@pytest.mark.asyncio
async def test_batch_file_mode_disabled_when_s3_v2_is_cold_storage_logger(monkeypatch: pytest.MonkeyPatch) -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
s3_aws_secret_access_key="test-secret",
s3_region_name="us-east-1",
s3_batch_file_upload=True,
)
put = _RecordingPut()
logger.async_httpx_client = AsyncMock()
logger.async_httpx_client.put = put
import litellm
monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2")
logger.log_queue = [_element({"id": "req-1"}, "0")]
await logger.async_send_batch()
assert len(put.calls) == 1
assert put.calls[0][0].endswith("test-0.json")
monkeypatch.setattr(litellm, "cold_storage_custom_logger", None)
logger.log_queue = [_element({"id": "req-2"}, "1")]
await logger.async_send_batch()
assert len(put.calls) == 2
assert put.calls[1][0].endswith(".jsonl")