From 09ebb28473e6e9e09c80ce2f822b88d8bc24f2e4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 07:32:00 +0000 Subject: [PATCH] fix(s3_v2): bound concurrent S3 uploads per flush and add opt-in JSONL batch files (#41258) * 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 --- litellm/constants.py | 1 + litellm/integrations/s3.py | 35 +- litellm/integrations/s3_v2.py | 126 +++- litellm/types/integrations/s3_v2.py | 2 + .../observability/_s3_v2_support.py | 344 ++++++++++ .../test_s3_v2_flush_surfaces.py | 97 +++ .../observability/test_s3_v2_upload_fanout.py | 630 ++++++++++++++++++ tests/test_litellm/integrations/test_s3_v2.py | 450 +++++++++++++ 8 files changed, 1665 insertions(+), 20 deletions(-) create mode 100644 tests/integration/observability/_s3_v2_support.py create mode 100644 tests/integration/observability/test_s3_v2_flush_surfaces.py create mode 100644 tests/integration/observability/test_s3_v2_upload_fanout.py diff --git a/litellm/constants.py b/litellm/constants.py index dcc12ef2ba9..3ff80d8b7dd 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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 diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index 54ec876fd0d..3c8619e82b2 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -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} diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 826f55cc798..dc33fe6c2bd 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -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", diff --git a/litellm/types/integrations/s3_v2.py b/litellm/types/integrations/s3_v2.py index 32864bf5b8c..555b16dc141 100644 --- a/litellm/types/integrations/s3_v2.py +++ b/litellm/types/integrations/s3_v2.py @@ -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" diff --git a/tests/integration/observability/_s3_v2_support.py b/tests/integration/observability/_s3_v2_support.py new file mode 100644 index 00000000000..104c0eda863 --- /dev/null +++ b/tests/integration/observability/_s3_v2_support.py @@ -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"SinkFailure", + 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) diff --git a/tests/integration/observability/test_s3_v2_flush_surfaces.py b/tests/integration/observability/test_s3_v2_flush_surfaces.py new file mode 100644 index 00000000000..2e0b7260a13 --- /dev/null +++ b/tests/integration/observability/test_s3_v2_flush_surfaces.py @@ -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" diff --git a/tests/integration/observability/test_s3_v2_upload_fanout.py b/tests/integration/observability/test_s3_v2_upload_fanout.py new file mode 100644 index 00000000000..3ebca152327 --- /dev/null +++ b/tests/integration/observability/test_s3_v2_upload_fanout.py @@ -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" diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index a9d13038180..c67eaa45112 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -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")