mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
fix(s3_v2): upload fresh events first, drop terminal failures and hour-old retries by default, opt-in adaptive concurrency (#43022)
* fix(s3_v2): drop terminal upload failures, bound retries per flush and enforce the queue cap at enqueue Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): keep retrying credential-rotation 403s, only AccessDenied-style errors are terminal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(env_keys): exclude DEFAULT_S3_MAX_FLUSH_ATTEMPTS as an internal tuning var Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): retry every 5xx, warn on first queue overflow, validate the flush budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): read the queue cap defensively so un-initialized loggers still enqueue Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): drop the getattr in _enqueue and tighten the retry tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): keep the constructor flush budget when the callback override is invalid Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(s3_v2): adapt per-object upload concurrency to sink latency and throttling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(s3_v2): tidy adaptive limiter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * perf(s3_v2): wake one waiter per released upload slot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(s3_v2): make the enqueue queue cap configurable with s3_max_queue_size Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): audit cells for cache hits, coded 403 and callback modes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): retry bucket-wide failures by default, age-budget requeues and make terminal drops and adaptive concurrency opt-in Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(s3_v2): suppress the missing-waiter ValueError explicitly in the adaptive limiter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): count oldest events trimmed after a failed flush as callback failures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(s3_v2): move the mutable-ok marker onto the list literal it suppresses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): report post-flush overflow drops once and grow adaptive concurrency above the floor before asserting back-off Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): fail the SlowDown back-off test when the measured window sees no PUTs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): restore base retry defaults, opt-in age budget, no enqueue cap, back off outside the limiter slot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): hoist the default no-op upload slot to a module constant Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): drop unused mutable-ok suppressions on queue appends Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): keep the retry queue oldest-first and prioritise fresh events at upload time Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(s3_v2): keep the mutable-ok marker on the queue list literal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): build request bodies inside the upload slot and keep the sync retry set at base parity Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): drop wall-clock sleeps from the unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): rebuild the request body inside the slot on every retry attempt Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): anchor the backoff window on the first observed failure and tighten shard assertions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): default the upload slot to the logger limiter so monkeypatched doubles keep working Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): clear ambient AWS env credentials so the rotating profile signs the sync retry test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): shrink the linear send-batch perf test to 2k/8k elements Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): fix stale batch sizes in the perf test assert message Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): keep async in-call retries on the base 403/500/503 set Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): hold the upload slot across retries, restore the bool upload contract, and fail safe on bool config Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): mark dropped uploads by element identity so a shared key cannot mask a retryable sibling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): make the per-flush drop lookup constant time Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): take the upload slot in the caller like base, build the body once per attempt loop Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): match base retry, logging and hook behaviour unless the new options are opted in Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): assert the wire key in the init-bypassed sync upload test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): assert the signed headers and wire key in the init-bypassed sync upload test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): default the upload limiter at class level instead of reading it with getattr Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): drop the duplicate annotations that redeclare the class-level counters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(s3_v2): drop terminal-failed uploads by default and bound retry age to one hour Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): fall back to the configured retry age on invalid values Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): drive retry-age tests from a fixed clock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): audit cells for retry-age opt-out and 429 single-put parity Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
3743c8563e
commit
e47b1f2a3f
12 changed files with 3083 additions and 144 deletions
|
|
@ -49,6 +49,7 @@ DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SE
|
|||
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"))
|
||||
DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY: Final = get_env_int("DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY", 200)
|
||||
# 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
|
||||
|
|
|
|||
78
litellm/integrations/adaptive_concurrency.py
Normal file
78
litellm/integrations/adaptive_concurrency.py
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
"""
|
||||
Adaptive in-flight concurrency limiter (AIMD, Vector ARC style).
|
||||
|
||||
Grows the limit additively after `limit` consecutive clean completions and
|
||||
halves it only on an explicit throttle signal (429, 503, SlowDown, or a
|
||||
transport error out of the PUT). With floor == ceiling it degenerates to a
|
||||
fixed-width semaphore.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections import deque
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PutSample:
|
||||
throttled: bool
|
||||
|
||||
|
||||
class AdaptiveConcurrencyLimiter:
|
||||
"""AIMD in-flight limiter used as `async with limiter:`."""
|
||||
|
||||
def __init__(self, initial: int, floor: int, ceiling: int) -> None:
|
||||
if not 1 <= floor <= ceiling:
|
||||
raise ValueError(f"adaptive limiter bounds must satisfy 1 <= floor <= ceiling, got {floor}..{ceiling}")
|
||||
self._limit: int = min(max(initial, floor), ceiling)
|
||||
self._floor: Final[int] = floor
|
||||
self._ceiling: Final[int] = ceiling
|
||||
self._clean_streak: int = 0
|
||||
self._in_flight: int = 0
|
||||
self._waiters: deque[asyncio.Future[None]] = deque() # mutable-ok: waiters queue up behind a full limit
|
||||
|
||||
@property
|
||||
def limit(self) -> int:
|
||||
return self._limit
|
||||
|
||||
async def __aenter__(self) -> "AdaptiveConcurrencyLimiter":
|
||||
if self._in_flight < self._limit:
|
||||
self._in_flight += 1
|
||||
return self
|
||||
waiter: Final = asyncio.get_running_loop().create_future()
|
||||
self._waiters.append(waiter)
|
||||
try:
|
||||
await waiter
|
||||
except asyncio.CancelledError:
|
||||
if waiter.done() and not waiter.cancelled():
|
||||
self._in_flight -= 1
|
||||
self._grant()
|
||||
else:
|
||||
with suppress(ValueError):
|
||||
self._waiters.remove(waiter)
|
||||
raise
|
||||
return self
|
||||
|
||||
def _grant(self) -> None:
|
||||
while self._in_flight < self._limit and self._waiters:
|
||||
waiter = self._waiters.popleft()
|
||||
if waiter.done():
|
||||
continue
|
||||
self._in_flight += 1
|
||||
waiter.set_result(None)
|
||||
|
||||
async def __aexit__(self, *_: object) -> None:
|
||||
self._in_flight -= 1
|
||||
self._grant()
|
||||
|
||||
def record(self, sample: PutSample) -> None:
|
||||
if sample.throttled:
|
||||
self._limit = max(self._floor, self._limit // 2)
|
||||
self._clean_streak = 0
|
||||
return
|
||||
self._clean_streak += 1
|
||||
if self._clean_streak >= self._limit and self._limit < self._ceiling:
|
||||
self._limit += 1
|
||||
self._clean_streak = 0
|
||||
self._grant()
|
||||
|
|
@ -36,22 +36,86 @@ def resolve_s3_log_prompts_only(configured: object, environ: Mapping[str, str] |
|
|||
return True
|
||||
|
||||
|
||||
def resolve_s3_max_concurrent_uploads(configured: object, fallback: int) -> int:
|
||||
def _resolve_positive_int(setting: str, configured: object, fallback: int, *, reject_bool: bool) -> int:
|
||||
if configured is None or configured == "":
|
||||
return fallback
|
||||
if reject_bool and isinstance(configured, bool):
|
||||
verbose_logger.warning(
|
||||
"s3 logging: %s=%r is a boolean, not an integer, using %s", setting, configured, fallback
|
||||
)
|
||||
return fallback
|
||||
try:
|
||||
bound: Final = _UPLOAD_BOUND.validate_python(configured.strip() if isinstance(configured, str) else configured)
|
||||
except ValidationError:
|
||||
verbose_logger.warning("s3 logging: %s=%r is not an integer, using %s", setting, configured, fallback)
|
||||
return fallback
|
||||
if bound < 1:
|
||||
verbose_logger.warning("s3 logging: %s=%r must be at least 1, using %s", setting, configured, fallback)
|
||||
return fallback
|
||||
return bound
|
||||
|
||||
|
||||
def resolve_s3_max_concurrent_uploads(configured: object, fallback: int) -> int:
|
||||
return _resolve_positive_int("s3_max_concurrent_uploads", configured, fallback, reject_bool=False)
|
||||
|
||||
|
||||
def resolve_s3_max_queue_size(configured: object, fallback: int) -> int:
|
||||
return _resolve_positive_int("s3_max_queue_size", configured, fallback, reject_bool=True)
|
||||
|
||||
|
||||
def resolve_s3_max_retry_age_seconds(configured: object, fallback: int | None) -> int | None:
|
||||
if configured is None or configured == "":
|
||||
return None
|
||||
if isinstance(configured, bool):
|
||||
verbose_logger.warning(
|
||||
"s3 logging: s3_max_retry_age_seconds=%r is a boolean, not an integer, falling back to %r",
|
||||
configured,
|
||||
fallback,
|
||||
)
|
||||
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
|
||||
"s3 logging: s3_max_retry_age_seconds=%r is not an integer, falling back to %r", configured, fallback
|
||||
)
|
||||
return fallback
|
||||
if bound < 1:
|
||||
if bound < 0:
|
||||
verbose_logger.warning(
|
||||
"s3 logging: s3_max_concurrent_uploads=%r must be at least 1, using %s", configured, fallback
|
||||
"s3 logging: s3_max_retry_age_seconds=%r must be at least 0, falling back to %r", configured, fallback
|
||||
)
|
||||
return fallback
|
||||
return bound
|
||||
return bound or None
|
||||
|
||||
|
||||
def resolve_s3_max_adaptive_concurrency(configured: object, fallback: int) -> int:
|
||||
return _resolve_positive_int("s3_max_adaptive_concurrency", configured, fallback, reject_bool=True)
|
||||
|
||||
|
||||
def resolve_s3_drop_on_terminal_error(configured: object) -> bool:
|
||||
if configured is None or configured == "":
|
||||
return True
|
||||
try:
|
||||
return _S3_BOOL.validate_python(configured.strip() if isinstance(configured, str) else configured)
|
||||
except ValidationError:
|
||||
verbose_logger.warning(
|
||||
"s3 logging: s3_drop_on_terminal_error=%r is not a boolean, dropping terminal-failed uploads",
|
||||
configured,
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
def resolve_s3_adaptive_concurrency(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_adaptive_concurrency=%r is not a boolean, keeping the fixed upload width",
|
||||
configured,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def resolve_s3_batch_file_upload(configured: object) -> bool:
|
||||
|
|
|
|||
|
|
@ -3,14 +3,19 @@ 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; 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
|
||||
NOTE 1: S3 does not provide a BATCH PUT API endpoint; by default each element is uploaded concurrently with the fixed s3_max_concurrent_uploads bound (or an adaptive bound when s3_adaptive_concurrency is on, backing off only on throttling), or with s3_batch_file_upload the whole flush is written as one .jsonl file
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Final, Literal, cast
|
||||
from urllib.parse import quote
|
||||
from uuid import uuid4
|
||||
|
||||
|
|
@ -21,15 +26,22 @@ from litellm._logging import print_verbose, verbose_logger
|
|||
from litellm.constants import (
|
||||
DEFAULT_S3_BATCH_SIZE,
|
||||
DEFAULT_S3_FLUSH_INTERVAL_SECONDS,
|
||||
DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY,
|
||||
DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
|
||||
)
|
||||
from litellm.integrations.adaptive_concurrency import AdaptiveConcurrencyLimiter, PutSample
|
||||
from litellm.integrations.s3 import (
|
||||
get_s3_object_download_filename,
|
||||
get_s3_object_key,
|
||||
prompts_only_payload,
|
||||
resolve_s3_adaptive_concurrency,
|
||||
resolve_s3_batch_file_upload,
|
||||
resolve_s3_drop_on_terminal_error,
|
||||
resolve_s3_log_prompts_only,
|
||||
resolve_s3_max_adaptive_concurrency,
|
||||
resolve_s3_max_concurrent_uploads,
|
||||
resolve_s3_max_queue_size,
|
||||
resolve_s3_max_retry_age_seconds,
|
||||
resolve_sse_params,
|
||||
)
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
|
|
@ -50,6 +62,42 @@ if TYPE_CHECKING:
|
|||
from botocore.credentials import Credentials
|
||||
|
||||
|
||||
UploadOutcome = Literal["delivered", "retry", "dropped"]
|
||||
|
||||
_TERMINAL_ERROR_CODES: Final = frozenset(
|
||||
{
|
||||
"EntityTooLarge",
|
||||
"InvalidArgument",
|
||||
"MalformedXML",
|
||||
"InvalidDigest",
|
||||
"KeyTooLongError",
|
||||
"BadDigest",
|
||||
"InvalidRequest",
|
||||
}
|
||||
)
|
||||
_BODY_CODED_STATUSES: Final = frozenset({400, 403})
|
||||
_RETRYABLE_STATUSES: Final = frozenset({403, 500, 503})
|
||||
_S3_ERROR_CODE: Final = re.compile(r"<Code>([^<]+)</Code>")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PreparedPut:
|
||||
json_string: str
|
||||
headers: Mapping[str, str]
|
||||
|
||||
|
||||
def _s3_error_code(response: httpx.Response) -> str | None:
|
||||
text: Final = response.text
|
||||
match: Final = _S3_ERROR_CODE.search(text) if isinstance(text, str) else None
|
||||
return match.group(1) if match else None
|
||||
|
||||
|
||||
def _is_terminal(response: httpx.Response) -> bool:
|
||||
"""True only for object-specific, unrecoverable rejections (400/403 with a terminal XML code).
|
||||
Unknown codes, empty or non-XML bodies, and every other status fail safe toward retry."""
|
||||
return response.status_code in _BODY_CODED_STATUSES and _s3_error_code(response) in _TERMINAL_ERROR_CODES
|
||||
|
||||
|
||||
def _s3_key_parent(s3_object_key: str) -> str:
|
||||
return s3_object_key.rsplit("/", 1)[0] if "/" in s3_object_key else ""
|
||||
|
||||
|
|
@ -58,11 +106,19 @@ 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")
|
||||
super().__init__(f"{failed} of {total} S3 uploads failed; transient failures kept in queue for the next flush")
|
||||
|
||||
|
||||
_in_flush: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar("s3_v2_in_flush", default=False)
|
||||
|
||||
|
||||
class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
||||
preserve_events_added_during_flush = True
|
||||
_flush_retries: int = 0
|
||||
_requeued_count: int = 0
|
||||
_upload_limiter: asyncio.Semaphore | AdaptiveConcurrencyLimiter | None = None
|
||||
s3_drop_on_terminal_error: bool = True
|
||||
s3_max_retry_age_seconds: int | None = 3600
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -92,6 +148,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
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_max_queue_size: int | None = None,
|
||||
s3_max_retry_age_seconds: int | None = 3600,
|
||||
s3_drop_on_terminal_error: bool = True,
|
||||
s3_adaptive_concurrency: bool = False,
|
||||
s3_max_adaptive_concurrency: int | None = None,
|
||||
s3_batch_file_upload: bool = False,
|
||||
s3_callback_params_override: dict | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -135,9 +196,22 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
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_max_queue_size=s3_max_queue_size,
|
||||
s3_max_retry_age_seconds=s3_max_retry_age_seconds,
|
||||
s3_drop_on_terminal_error=s3_drop_on_terminal_error,
|
||||
s3_adaptive_concurrency=s3_adaptive_concurrency,
|
||||
s3_max_adaptive_concurrency=s3_max_adaptive_concurrency,
|
||||
s3_batch_file_upload=s3_batch_file_upload,
|
||||
)
|
||||
self._upload_semaphore = asyncio.Semaphore(self.s3_max_concurrent_uploads)
|
||||
self._upload_limiter = (
|
||||
AdaptiveConcurrencyLimiter(
|
||||
initial=self.s3_max_concurrent_uploads,
|
||||
floor=self.s3_max_concurrent_uploads,
|
||||
ceiling=max(self.s3_max_concurrent_uploads, self.s3_max_adaptive_concurrency),
|
||||
)
|
||||
if self.s3_adaptive_concurrency
|
||||
else asyncio.Semaphore(self.s3_max_concurrent_uploads)
|
||||
)
|
||||
verbose_logger.debug("s3 logger using endpoint url %s", s3_endpoint_url)
|
||||
|
||||
# IMPORTANT
|
||||
|
|
@ -158,8 +232,12 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
flush_lock=self.flush_lock,
|
||||
flush_interval=s3_flush_interval,
|
||||
batch_size=s3_batch_size,
|
||||
max_queue_size=self.s3_max_queue_size,
|
||||
)
|
||||
self.log_queue: list[s3BatchLoggingElement] = []
|
||||
self._requeued_count = 0
|
||||
self._flush_retries = 0
|
||||
self._flush_dropped: dict[int, s3BatchLoggingElement] = {}
|
||||
|
||||
# Call BaseAWSLLM's __init__
|
||||
BaseAWSLLM.__init__(self)
|
||||
|
|
@ -194,6 +272,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
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_max_queue_size: int | None = None,
|
||||
s3_max_retry_age_seconds: int | None = 3600,
|
||||
s3_drop_on_terminal_error: bool = True,
|
||||
s3_adaptive_concurrency: bool = False,
|
||||
s3_max_adaptive_concurrency: int | None = None,
|
||||
s3_batch_file_upload: bool = False,
|
||||
params_source: dict | None = None,
|
||||
):
|
||||
|
|
@ -259,6 +342,37 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
|
||||
)
|
||||
|
||||
configured_queue_size: Final = params.get("s3_max_queue_size")
|
||||
constructor_queue_size: Final = resolve_s3_max_queue_size(
|
||||
s3_max_queue_size, CustomBatchLogger.DEFAULT_MAX_QUEUE_SIZE
|
||||
)
|
||||
self.s3_max_queue_size = resolve_s3_max_queue_size(configured_queue_size, constructor_queue_size)
|
||||
|
||||
configured_retry_age: Final = params.get("s3_max_retry_age_seconds")
|
||||
constructor_retry_age: Final = resolve_s3_max_retry_age_seconds(s3_max_retry_age_seconds, 3600)
|
||||
self.s3_max_retry_age_seconds = (
|
||||
constructor_retry_age
|
||||
if configured_retry_age is None or configured_retry_age == ""
|
||||
else resolve_s3_max_retry_age_seconds(configured_retry_age, constructor_retry_age)
|
||||
)
|
||||
|
||||
configured_drop: Final = params.get("s3_drop_on_terminal_error")
|
||||
self.s3_drop_on_terminal_error = resolve_s3_drop_on_terminal_error(
|
||||
configured_drop if configured_drop is not None else s3_drop_on_terminal_error
|
||||
)
|
||||
|
||||
self.s3_adaptive_concurrency = s3_adaptive_concurrency or resolve_s3_adaptive_concurrency(
|
||||
params.get("s3_adaptive_concurrency")
|
||||
)
|
||||
|
||||
configured_adaptive_ceiling: Final = params.get("s3_max_adaptive_concurrency")
|
||||
self.s3_max_adaptive_concurrency = resolve_s3_max_adaptive_concurrency(
|
||||
s3_max_adaptive_concurrency
|
||||
if configured_adaptive_ceiling is None or configured_adaptive_ceiling == ""
|
||||
else configured_adaptive_ceiling,
|
||||
DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY,
|
||||
)
|
||||
|
||||
self.s3_batch_file_upload = s3_batch_file_upload or resolve_s3_batch_file_upload(
|
||||
params.get("s3_batch_file_upload")
|
||||
)
|
||||
|
|
@ -310,6 +424,35 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
}
|
||||
return {key: value for key, value in candidates.items() if value}
|
||||
|
||||
def _prepare_put(self, batch_logging_element: s3BatchLoggingElement) -> _PreparedPut:
|
||||
try:
|
||||
import base64
|
||||
import hashlib
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
|
||||
json_string: Final = (
|
||||
batch_logging_element.body
|
||||
if batch_logging_element.body is not None
|
||||
else safe_dumps(batch_logging_element.payload)
|
||||
)
|
||||
content_hash: Final = hashlib.sha256(json_string.encode("utf-8")).hexdigest()
|
||||
content_md5: Final = base64.b64encode(
|
||||
hashlib.md5(json_string.encode("utf-8"), usedforsecurity=False).digest()
|
||||
).decode()
|
||||
return _PreparedPut(
|
||||
json_string=json_string,
|
||||
headers={
|
||||
"Content-Type": batch_logging_element.content_type,
|
||||
"Content-MD5": content_md5,
|
||||
"x-amz-content-sha256": content_hash,
|
||||
"Content-Language": "en",
|
||||
"Content-Disposition": f'inline; filename="{batch_logging_element.s3_object_download_filename}"',
|
||||
"Cache-Control": "private, immutable, max-age=31536000, s-maxage=0",
|
||||
**self._sse_headers(),
|
||||
},
|
||||
)
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
await self._async_log_event_base(
|
||||
kwargs=kwargs,
|
||||
|
|
@ -384,12 +527,21 @@ 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) -> bool:
|
||||
try:
|
||||
import base64
|
||||
import hashlib
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
@property
|
||||
def _upload_semaphore(self) -> asyncio.Semaphore | AdaptiveConcurrencyLimiter:
|
||||
limiter: Final = self._upload_limiter
|
||||
if limiter is None:
|
||||
raise AttributeError("_upload_semaphore")
|
||||
return limiter
|
||||
|
||||
@_upload_semaphore.setter
|
||||
def _upload_semaphore(self, value: asyncio.Semaphore | AdaptiveConcurrencyLimiter) -> None:
|
||||
self._upload_limiter = value
|
||||
|
||||
async def async_upload_data_to_s3(
|
||||
self,
|
||||
batch_logging_element: s3BatchLoggingElement,
|
||||
) -> bool:
|
||||
try:
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
|
||||
|
|
@ -400,31 +552,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
|
||||
url: Final = self._build_object_url(batch_logging_element.s3_object_key)
|
||||
|
||||
# Convert JSON to string
|
||||
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()
|
||||
content_md5: Final = base64.b64encode(
|
||||
hashlib.md5(json_string.encode("utf-8"), usedforsecurity=False).digest()
|
||||
).decode()
|
||||
|
||||
# Prepare the request
|
||||
headers: Final = {
|
||||
"Content-Type": batch_logging_element.content_type,
|
||||
"Content-MD5": content_md5,
|
||||
"x-amz-content-sha256": content_hash,
|
||||
"Content-Language": "en",
|
||||
"Content-Disposition": f'inline; filename="{batch_logging_element.s3_object_download_filename}"',
|
||||
"Cache-Control": "private, immutable, max-age=31536000, s-maxage=0",
|
||||
**self._sse_headers(),
|
||||
}
|
||||
|
||||
async def signed_put() -> httpx.Response:
|
||||
async def signed_put(prepared: _PreparedPut) -> httpx.Response:
|
||||
credentials: Final = await asyncified_get_credentials(
|
||||
aws_access_key_id=self.s3_aws_access_key_id,
|
||||
aws_secret_access_key=self.s3_aws_secret_access_key,
|
||||
|
|
@ -436,18 +564,26 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
aws_web_identity_token=self.s3_aws_web_identity_token,
|
||||
aws_sts_endpoint=self.s3_aws_sts_endpoint,
|
||||
)
|
||||
signed_headers: Final = await run_aws_signing(self._sign_put, credentials, url, json_string, headers)
|
||||
signed_headers: Final = await run_aws_signing(
|
||||
self._sign_put, credentials, url, prepared.json_string, prepared.headers
|
||||
)
|
||||
try:
|
||||
return await self.async_httpx_client.put(url, data=json_string, headers=signed_headers)
|
||||
return await self.async_httpx_client.put(url, data=prepared.json_string, headers=signed_headers)
|
||||
except httpx.HTTPStatusError as error:
|
||||
return error.response
|
||||
|
||||
max_retries: Final = 3
|
||||
prepared: Final = self._prepare_put(batch_logging_element)
|
||||
for attempt in range(max_retries):
|
||||
response = await signed_put()
|
||||
if response.status_code in (403, 500, 503) and attempt < max_retries - 1:
|
||||
response = await self._recorded_put(partial(signed_put, prepared))
|
||||
if (
|
||||
response.status_code in _RETRYABLE_STATUSES
|
||||
and not (self.s3_drop_on_terminal_error and _is_terminal(response))
|
||||
and attempt < max_retries - 1
|
||||
):
|
||||
wait_time = 2**attempt # 1s, 2s
|
||||
verbose_logger.warning(
|
||||
verbose_logger.log(
|
||||
logging.DEBUG if _in_flush.get() else logging.WARNING,
|
||||
"S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s",
|
||||
response.status_code,
|
||||
wait_time,
|
||||
|
|
@ -455,6 +591,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
max_retries,
|
||||
batch_logging_element.s3_object_key,
|
||||
)
|
||||
self._flush_retries += 1
|
||||
await asyncio.sleep(wait_time)
|
||||
continue
|
||||
response.raise_for_status()
|
||||
|
|
@ -462,6 +599,13 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
except Exception as e:
|
||||
verbose_logger.exception("Error uploading to s3: %s", e)
|
||||
self.handle_callback_failure(callback_name="S3Logger")
|
||||
if isinstance(e, httpx.HTTPStatusError) and self.s3_drop_on_terminal_error and _is_terminal(e.response):
|
||||
verbose_logger.warning(
|
||||
"s3 logging: dropping object %s after terminal status %s",
|
||||
batch_logging_element.s3_object_key,
|
||||
e.response.status_code,
|
||||
)
|
||||
self._flush_dropped[id(batch_logging_element)] = batch_logging_element
|
||||
return False
|
||||
return True
|
||||
|
||||
|
|
@ -483,11 +627,64 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
# see custom_batch_logger.py which triggers the flush
|
||||
#########################################################
|
||||
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:
|
||||
self._flush_retries = 0
|
||||
self._flush_dropped = {} # mutable-ok: per-flush drop marks read back by _upload_bounded
|
||||
stale: Final = min(self._requeued_count, len(uploads)) if len(uploads) == len(batch) else 0
|
||||
order: Final = (*range(stale, len(uploads)), *range(stale))
|
||||
ordered: Final = await asyncio.gather(*(self._upload_outcome(uploads[i]) for i in order))
|
||||
outcomes: Final = dict(zip(order, ordered, strict=True))
|
||||
results: Final = tuple(outcomes[i] for i in range(len(uploads)))
|
||||
if self._flush_retries:
|
||||
verbose_logger.warning(
|
||||
"s3 logging: %s in-call retries across %s uploads this flush",
|
||||
self._flush_retries,
|
||||
len(uploads),
|
||||
)
|
||||
delivered: Final = sum(1 for outcome in results if outcome == "delivered")
|
||||
bucket_wide: Final = delivered == 0
|
||||
failed: Final = tuple(
|
||||
(element, outcome) for element, outcome in zip(uploads, results, strict=True) if outcome != "delivered"
|
||||
)
|
||||
now: Final = time.monotonic()
|
||||
requeued: Final = (
|
||||
tuple(element for element, _ in failed)
|
||||
if bucket_wide
|
||||
else tuple(
|
||||
element
|
||||
if element.retrying_since is not None or self.s3_max_retry_age_seconds is None
|
||||
else element.model_copy(update={"retrying_since": now})
|
||||
for element, outcome in failed
|
||||
if outcome != "dropped"
|
||||
and not (
|
||||
self.s3_max_retry_age_seconds is not None
|
||||
and element.retrying_since is not None
|
||||
and now - element.retrying_since > self.s3_max_retry_age_seconds
|
||||
)
|
||||
)
|
||||
)
|
||||
dropped: Final = len(failed) - len(requeued)
|
||||
if dropped:
|
||||
verbose_logger.warning(
|
||||
"s3 logging: %s uploads dropped (terminal or retrying longer than s3_max_retry_age_seconds=%s)",
|
||||
dropped,
|
||||
self.s3_max_retry_age_seconds,
|
||||
)
|
||||
if not requeued:
|
||||
self._requeued_count = 0
|
||||
return
|
||||
self.log_queue = [*failed, *self.log_queue[len(batch) :]]
|
||||
arrivals: Final = self.log_queue[len(batch) :]
|
||||
overflow: Final = max(0, len(requeued) + len(arrivals) - self.max_queue_size)
|
||||
if overflow:
|
||||
verbose_logger.warning(
|
||||
"s3 logging: queue exceeded max_queue_size=%s after a failed flush, dropped %s oldest events",
|
||||
self.max_queue_size,
|
||||
overflow,
|
||||
)
|
||||
self.log_queue = [ # mutable-ok: log_queue is the flush buffer shared with custom_batch_logger
|
||||
*requeued,
|
||||
*arrivals,
|
||||
][overflow:]
|
||||
self._requeued_count = max(0, len(requeued) - overflow)
|
||||
raise S3BatchUploadError(failed=len(failed), total=len(uploads))
|
||||
|
||||
def _batch_file_mode_active(self) -> bool:
|
||||
|
|
@ -502,8 +699,37 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
return True
|
||||
|
||||
async def _upload_bounded(self, element: s3BatchLoggingElement) -> bool:
|
||||
async with self._upload_semaphore:
|
||||
return await self.async_upload_data_to_s3(element)
|
||||
token: Final = _in_flush.set(True)
|
||||
try:
|
||||
async with self._upload_semaphore:
|
||||
return await self.async_upload_data_to_s3(element)
|
||||
finally:
|
||||
_in_flush.reset(token)
|
||||
|
||||
async def _upload_outcome(self, element: s3BatchLoggingElement) -> UploadOutcome:
|
||||
delivered: Final = await self._upload_bounded(element)
|
||||
if delivered:
|
||||
return "delivered"
|
||||
if id(element) in self._flush_dropped:
|
||||
return "dropped"
|
||||
return "retry"
|
||||
|
||||
async def _recorded_put(self, signed_put: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response:
|
||||
limiter: Final = self._upload_limiter
|
||||
adaptive: Final = limiter if isinstance(limiter, AdaptiveConcurrencyLimiter) else None
|
||||
try:
|
||||
response: Final = await signed_put()
|
||||
except Exception:
|
||||
if adaptive is not None:
|
||||
adaptive.record(PutSample(throttled=True))
|
||||
raise
|
||||
if adaptive is not None:
|
||||
adaptive.record(
|
||||
PutSample(
|
||||
throttled=response.status_code in (429, 503) or _s3_error_code(response) == "SlowDown",
|
||||
)
|
||||
)
|
||||
return response
|
||||
|
||||
def _batch_file_elements(self, batch: tuple[s3BatchLoggingElement, ...]) -> tuple[s3BatchLoggingElement, ...]:
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
|
|
@ -527,6 +753,9 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
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",
|
||||
retrying_since=min(
|
||||
(element.retrying_since for element in elements if element.retrying_since is not None), default=None
|
||||
),
|
||||
)
|
||||
|
||||
def create_s3_batch_logging_element(
|
||||
|
|
@ -596,58 +825,35 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
)
|
||||
|
||||
def upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement):
|
||||
try:
|
||||
import base64
|
||||
import hashlib
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
try:
|
||||
verbose_logger.debug("s3_v2 logger - uploading data to s3 - %s", batch_logging_element.s3_object_key)
|
||||
|
||||
url: Final = self._build_object_url(batch_logging_element.s3_object_key)
|
||||
|
||||
# Convert JSON to string
|
||||
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()
|
||||
content_md5: Final = base64.b64encode(
|
||||
hashlib.md5(json_string.encode("utf-8"), usedforsecurity=False).digest()
|
||||
).decode()
|
||||
|
||||
# Prepare the request
|
||||
headers: Final = {
|
||||
"Content-Type": batch_logging_element.content_type,
|
||||
"Content-MD5": content_md5,
|
||||
"x-amz-content-sha256": content_hash,
|
||||
"Content-Language": "en",
|
||||
"Content-Disposition": f'inline; filename="{batch_logging_element.s3_object_download_filename}"',
|
||||
"Cache-Control": "private, immutable, max-age=31536000, s-maxage=0",
|
||||
**self._sse_headers(),
|
||||
}
|
||||
prepared: Final = self._prepare_put(batch_logging_element)
|
||||
|
||||
httpx_client: Final = _get_httpx_client(
|
||||
params=({"ssl_verify": self.s3_verify} if self.s3_verify is not None else None)
|
||||
)
|
||||
|
||||
def signed_put() -> httpx.Response:
|
||||
def signed_put(prepared_put: _PreparedPut) -> httpx.Response:
|
||||
credentials: Final = self.get_credentials(
|
||||
aws_access_key_id=self.s3_aws_access_key_id,
|
||||
aws_secret_access_key=self.s3_aws_secret_access_key,
|
||||
aws_session_token=self.s3_aws_session_token,
|
||||
aws_region_name=self.s3_region_name,
|
||||
)
|
||||
signed_headers: Final = self._sign_put(credentials, url, json_string, headers)
|
||||
return httpx_client.put(url, data=json_string, headers=signed_headers)
|
||||
signed_headers: Final = self._sign_put(credentials, url, prepared_put.json_string, prepared_put.headers)
|
||||
return httpx_client.put(url, data=prepared_put.json_string, headers=signed_headers)
|
||||
|
||||
max_retries: Final = 3
|
||||
for attempt in range(max_retries):
|
||||
response = signed_put()
|
||||
if response.status_code in (403, 500, 503) and attempt < max_retries - 1:
|
||||
response = signed_put(prepared)
|
||||
if (
|
||||
response.status_code in _RETRYABLE_STATUSES
|
||||
and not (self.s3_drop_on_terminal_error and _is_terminal(response))
|
||||
and attempt < max_retries - 1
|
||||
):
|
||||
wait_time = 2**attempt # 1s, 2s
|
||||
verbose_logger.warning(
|
||||
"S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s",
|
||||
|
|
@ -664,6 +870,12 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
except Exception as e:
|
||||
verbose_logger.exception("Error uploading to s3: %s", e)
|
||||
self.handle_callback_failure(callback_name="S3Logger")
|
||||
if isinstance(e, httpx.HTTPStatusError) and self.s3_drop_on_terminal_error and _is_terminal(e.response):
|
||||
verbose_logger.warning(
|
||||
"s3 logging: dropping object %s after terminal status %s",
|
||||
batch_logging_element.s3_object_key,
|
||||
e.response.status_code,
|
||||
)
|
||||
|
||||
async def _download_object_from_s3(self, s3_object_key: str) -> dict | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -11,3 +11,4 @@ class s3BatchLoggingElement(BaseModel):
|
|||
s3_object_download_filename: str
|
||||
body: str | None = None
|
||||
content_type: str = "application/json"
|
||||
retrying_since: float | None = None
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ EXCLUDED_ROLLOUT_FLAGS = {
|
|||
EXCLUDED_INTERNAL_TUNING_VARS = {
|
||||
"ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS",
|
||||
"ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE",
|
||||
"DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY",
|
||||
}
|
||||
|
||||
EXCLUDED_TERMINAL_VARS = {
|
||||
|
|
|
|||
|
|
@ -27,11 +27,15 @@ class RecordingS3Sink:
|
|||
fail_attempts: int = 0
|
||||
fail_until: float = 0.0
|
||||
fail_status: int = 503
|
||||
fail_code: str = "SinkFailure"
|
||||
fail_body: bytes | None = None
|
||||
delay_seconds: float = 0.5
|
||||
lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
in_flight: int = 0
|
||||
peak: int = 0
|
||||
attempts: int = 0
|
||||
attempt_log: list[tuple[float, int]] = field(default_factory=list) # mutable-ok: appended under lock per PUT
|
||||
attempt_counts: dict[str, int] = field(default_factory=dict) # mutable-ok: per-target PUT counts under lock
|
||||
store: dict[str, bytes] = field(default_factory=dict) # mutable-ok: GET reads must see writes from earlier PUTs
|
||||
|
||||
def respond(self, request: Request) -> Reply:
|
||||
|
|
@ -44,20 +48,31 @@ class RecordingS3Sink:
|
|||
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.attempt_counts[request.target] = self.attempt_counts.get(request.target, 0) + 1
|
||||
self.in_flight += 1
|
||||
self.peak = max(self.peak, self.in_flight)
|
||||
self.store[request.target] = request.body
|
||||
self.attempt_log.append((time.time(), self.in_flight))
|
||||
failing: Final = self.attempts <= self.fail_attempts or time.time() < self.fail_until
|
||||
if not failing:
|
||||
self.store[request.target] = request.body
|
||||
time.sleep(self.delay_seconds)
|
||||
with self.lock:
|
||||
self.in_flight -= 1
|
||||
if failing:
|
||||
return Reply(
|
||||
status=self.fail_status,
|
||||
body=self.fail_body
|
||||
if self.fail_body is not None
|
||||
else f"<Error><Code>{self.fail_code}</Code></Error>".encode(),
|
||||
content_type="application/xml",
|
||||
)
|
||||
return Reply()
|
||||
|
||||
def peak_between(self, start: float, end: float) -> int:
|
||||
with self.lock:
|
||||
samples: Final = tuple(in_flight for when, in_flight in self.attempt_log if start <= when < end)
|
||||
return max(samples, default=0)
|
||||
|
||||
def objects(self) -> Mapping[str, bytes]:
|
||||
with self.lock:
|
||||
return MappingProxyType(dict(self.store))
|
||||
|
|
@ -240,7 +255,13 @@ SURFACES: Final = ("chat", "chat_stream", "messages", "messages_stream", "respon
|
|||
|
||||
|
||||
def call_surface(
|
||||
candidate: Gateway, surface: str, openai_model: str, anthropic_model: str, key: str, marker: str
|
||||
candidate: Gateway,
|
||||
surface: str,
|
||||
openai_model: str,
|
||||
anthropic_model: str,
|
||||
key: str,
|
||||
marker: str,
|
||||
no_cache: bool = True,
|
||||
) -> 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("/")
|
||||
|
|
@ -249,7 +270,7 @@ def call_surface(
|
|||
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}},
|
||||
extra_body={"cache": {"no-cache": True}} if no_cache else {},
|
||||
)
|
||||
return reply.id, None
|
||||
|
||||
|
|
@ -258,7 +279,7 @@ def call_surface(
|
|||
model=openai_model,
|
||||
messages=[{"role": "user", "content": marker}],
|
||||
stream=True,
|
||||
extra_body={"cache": {"no-cache": True}},
|
||||
extra_body={"cache": {"no-cache": True}} if no_cache else {},
|
||||
)
|
||||
seen = ""
|
||||
async for chunk in stream:
|
||||
|
|
@ -283,7 +304,7 @@ def call_surface(
|
|||
response: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": openai_model, "input": marker, "cache": {"no-cache": True}},
|
||||
{"model": openai_model, "input": marker, **({"cache": {"no-cache": True}} if no_cache else {})},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
|
@ -339,6 +360,10 @@ def matched_ids(
|
|||
if payload["id"] in response_ids:
|
||||
landed.append(payload["id"])
|
||||
continue
|
||||
uncached: Final = str(payload["id"]).rsplit("_cache_hit", 1)[0]
|
||||
if uncached in response_ids:
|
||||
landed.append(str(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)
|
||||
|
|
|
|||
|
|
@ -1,20 +1,24 @@
|
|||
import os
|
||||
import re
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from redis import Redis
|
||||
from _s3_v2_support import (
|
||||
BUCKET,
|
||||
PREFIX,
|
||||
SURFACES,
|
||||
RecordingS3Sink,
|
||||
call_surface,
|
||||
collect_payloads,
|
||||
matched_ids,
|
||||
mixed_burst,
|
||||
s3_config,
|
||||
surface_reply,
|
||||
)
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import wire_server
|
||||
|
||||
|
|
@ -95,3 +99,38 @@ def test_s3_v2_sink_outage_mid_mixed_burst_recovers_every_response_id(gateway: G
|
|||
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"
|
||||
|
||||
|
||||
def test_s3_v2_cache_hit_twins_log_one_object_per_request(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3cache" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RecordingS3Sink(delay_seconds=0.1)
|
||||
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])
|
||||
cache: Final = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
|
||||
keys_before: Final = cache.dbsize()
|
||||
warmed: Final = tuple(
|
||||
call_surface(candidate, surface, openai_model, anthropic_model, key, f"{marker}-{surface}")
|
||||
for surface in SURFACES
|
||||
)
|
||||
eventually(cache.dbsize, lambda size: size >= keys_before + len(SURFACES), seconds=30)
|
||||
repeated: Final = tuple(
|
||||
call_surface(candidate, surface, openai_model, anthropic_model, key, f"{marker}-{surface}", False)
|
||||
for surface in SURFACES
|
||||
)
|
||||
payloads: Final = collect_payloads(sink, 2 * len(SURFACES))
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == len(SURFACES), (
|
||||
"a repeated request reached the upstream; the six repeats must all be served from cache"
|
||||
)
|
||||
assert len(payloads) == 12
|
||||
assert sum(1 for payload in payloads if payload["cache_hit"] is True) == 6
|
||||
assert sum(1 for payload in payloads if payload["cache_hit"] is not True) == 6
|
||||
assert matched_ids(payloads, warmed + repeated)
|
||||
|
|
|
|||
|
|
@ -119,8 +119,8 @@ BATCH_KEY: Final = re.compile(
|
|||
)
|
||||
|
||||
|
||||
@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:
|
||||
@pytest.mark.covers("other.observability.s3_v2.flush_bounds_concurrent_puts_to_default_ceiling_and_keeps_every_log")
|
||||
def test_s3_v2_flush_bounds_concurrent_puts_to_the_default_ceiling(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:
|
||||
|
|
@ -134,7 +134,9 @@ def test_s3_v2_flush_bounds_concurrent_puts_to_the_default_of_sixteen(gateway: G
|
|||
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 sink.peak <= 16, (
|
||||
f"peak concurrent PUTs {sink.peak} exceeded the default width of 16 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
|
||||
|
|
@ -272,7 +274,7 @@ def test_s3_v2_upstream_failure_events_land_alongside_successes(gateway: Gateway
|
|||
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.covers("other.observability.s3_v2.invalid_or_empty_bound_falls_back_to_default_ceiling")
|
||||
@pytest.mark.parametrize(
|
||||
("bad", "warns"),
|
||||
[
|
||||
|
|
@ -281,7 +283,7 @@ def test_s3_v2_upstream_failure_events_land_alongside_successes(gateway: Gateway
|
|||
pytest.param("", False, id="empty"),
|
||||
],
|
||||
)
|
||||
def test_s3_v2_invalid_or_empty_bound_falls_back_to_sixteen(
|
||||
def test_s3_v2_invalid_or_empty_bound_falls_back_to_default_ceiling(
|
||||
gateway: Gateway, tmp_path: Path, bad: JsonValue, warns: bool
|
||||
) -> None:
|
||||
marker: Final = "s3bound" + uuid.uuid4().hex[:8]
|
||||
|
|
@ -305,14 +307,14 @@ def test_s3_v2_invalid_or_empty_bound_falls_back_to_sixteen(
|
|||
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 sink.peak <= 16, f"peak concurrent PUTs {sink.peak} exceeded the fallback width of 16"
|
||||
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)
|
||||
sink: Final = RecordingS3Sink(fail_status=503, 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 (
|
||||
|
|
@ -336,6 +338,283 @@ def test_s3_v2_sink_rejection_requeues_and_delivers_every_id_once(gateway: Gatew
|
|||
assert frozenset(payload["id"] for payload in payloads) == ids
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class RejectingS3Sink:
|
||||
"""Answers every PUT whose body carries `reject_marker` with `reject_status`, accepts the rest,
|
||||
and counts the rejected attempts so a test can see whether the proxy keeps re-sending them."""
|
||||
|
||||
reject_marker: str
|
||||
reject_status: int
|
||||
reject_code: str = "AccessDenied"
|
||||
reject_until: float = float("inf")
|
||||
lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
rejected_attempts: int = 0
|
||||
rejected_times: list[float] = field(default_factory=list) # mutable-ok: appended under lock per rejected PUT
|
||||
store: dict[str, bytes] = field(default_factory=dict) # mutable-ok: later PUTs must be visible to earlier polls
|
||||
|
||||
def respond(self, request: Request) -> Reply:
|
||||
assert request.method == "PUT", request.method
|
||||
with self.lock:
|
||||
if self.reject_marker.encode() in request.body and time.time() < self.reject_until:
|
||||
self.rejected_attempts += 1
|
||||
self.rejected_times.append(time.time())
|
||||
return Reply(status=self.reject_status, body=f"<Error><Code>{self.reject_code}</Code></Error>".encode())
|
||||
self.store[request.target] = request.body
|
||||
return Reply()
|
||||
|
||||
def landed_ids(self) -> frozenset[str]:
|
||||
with self.lock:
|
||||
bodies: Final = tuple(self.store.values())
|
||||
return frozenset(json.loads(line)["id"] for body in bodies for line in body.splitlines())
|
||||
|
||||
|
||||
def _send(candidate: Gateway, model: str, key: str, identity: str) -> None:
|
||||
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
|
||||
|
||||
|
||||
def _send_and_wait_until_landed(candidate: Gateway, model: str, key: str, sink: RejectingS3Sink, identity: str) -> None:
|
||||
_send(candidate, model, key, identity)
|
||||
eventually(sink.landed_ids, lambda landed: identity in landed, seconds=60)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status", "code"),
|
||||
[
|
||||
pytest.param(403, "AccessDenied", id="access_denied"),
|
||||
pytest.param(404, "NoSuchBucket", id="no_such_bucket"),
|
||||
pytest.param(400, "KMS.DisabledException", id="kms_disabled"),
|
||||
],
|
||||
)
|
||||
def test_s3_v2_object_rejected_with_a_bucket_wide_code_is_delivered_once_the_fault_clears(
|
||||
gateway: Gateway, tmp_path: Path, status: int, code: str
|
||||
) -> None:
|
||||
marker: Final = "s3fault" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RejectingS3Sink(reject_marker=f"{marker}-denied", reject_status=status, reject_code=code)
|
||||
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": False})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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])
|
||||
_send(candidate, model, key, f"{marker}-denied")
|
||||
_send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-first-flush")
|
||||
eventually(lambda: sink.rejected_attempts, lambda attempts: attempts >= 2, seconds=30)
|
||||
sink.reject_until = time.time()
|
||||
eventually(sink.landed_ids, lambda landed: f"{marker}-denied" in landed, seconds=60)
|
||||
readiness: Final = candidate.client.get("/health/readiness")
|
||||
assert readiness.status_code == 200, readiness.text
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == 2
|
||||
assert sink.landed_ids() == {f"{marker}-denied", f"{marker}-first-flush"}
|
||||
|
||||
|
||||
def test_s3_v2_terminal_object_is_put_once_and_dropped_by_default(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3toolarge" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RejectingS3Sink(reject_marker=f"{marker}-huge", reject_status=400, reject_code="EntityTooLarge")
|
||||
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": False})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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])
|
||||
_send(candidate, model, key, f"{marker}-huge")
|
||||
_send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-sibling")
|
||||
_send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-second-flush")
|
||||
_send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-third-flush")
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == 4
|
||||
assert sink.rejected_attempts == 1, (
|
||||
f"an EntityTooLarge object was PUT {sink.rejected_attempts} times next to delivered siblings; "
|
||||
"with the default s3_drop_on_terminal_error it must be attempted once and dropped"
|
||||
)
|
||||
|
||||
|
||||
def test_s3_v2_terminal_object_keeps_retrying_when_opted_out(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3keep" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RejectingS3Sink(reject_marker=f"{marker}-huge", reject_status=400, reject_code="EntityTooLarge")
|
||||
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": False, "s3_drop_on_terminal_error": False}
|
||||
)
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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])
|
||||
_send(candidate, model, key, f"{marker}-huge")
|
||||
_send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-sibling")
|
||||
_send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-second-flush")
|
||||
eventually(lambda: sink.rejected_attempts, lambda attempts: attempts >= 2, seconds=30)
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == 3
|
||||
|
||||
|
||||
def test_s3_v2_aged_out_object_is_dropped_next_to_delivered_siblings(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3aged" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RejectingS3Sink(reject_marker=f"{marker}-doomed", reject_status=503, reject_code="InternalError")
|
||||
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": False, "s3_max_retry_age_seconds": 1})
|
||||
with (
|
||||
owned_proxy_process(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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])
|
||||
_send(owned.gateway, model, key, f"{marker}-doomed")
|
||||
_send_and_wait_until_landed(owned.gateway, model, key, sink, f"{marker}-sibling")
|
||||
_send(owned.gateway, model, key, f"{marker}-trigger")
|
||||
eventually(
|
||||
lambda: owned.log.read_text(),
|
||||
lambda text: "retrying longer than s3_max_retry_age_seconds=1)" in text,
|
||||
seconds=60,
|
||||
)
|
||||
exhausted: Final = sink.rejected_attempts
|
||||
_send_and_wait_until_landed(owned.gateway, model, key, sink, f"{marker}-one-flush-later")
|
||||
_send_and_wait_until_landed(owned.gateway, model, key, sink, f"{marker}-two-flushes-later")
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == 5
|
||||
assert 3 <= exhausted <= 3 * 3, f"{exhausted} PUTs for an object that aged out after its second flush"
|
||||
assert sink.rejected_attempts == exhausted, (
|
||||
f"a 503 object kept being PUT after ageing out: {exhausted} -> {sink.rejected_attempts}"
|
||||
)
|
||||
|
||||
|
||||
def test_s3_v2_aged_out_object_stays_queued_while_the_whole_sink_is_down(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3down" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RecordingS3Sink(fail_status=503, delay_seconds=0.1)
|
||||
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": False, "s3_max_retry_age_seconds": 1})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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() + 12
|
||||
ids: Final = _push(candidate, model, key, marker, 4)
|
||||
payloads: Final = collect_payloads(sink, 4, seconds=90)
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == 4
|
||||
assert frozenset(payload["id"] for payload in payloads) == ids, (
|
||||
"a bucket-wide outage longer than the age budget lost events"
|
||||
)
|
||||
|
||||
|
||||
def test_s3_v2_failing_sink_trims_the_oldest_failed_events_past_the_queue_cap(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3cap" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RecordingS3Sink(fail_status=503, delay_seconds=0.05)
|
||||
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": False, "s3_max_queue_size": 4})
|
||||
with (
|
||||
owned_proxy_process(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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() + 15
|
||||
_send(owned.gateway, model, key, f"{marker}-probe")
|
||||
eventually(lambda: owned.log.read_text(), lambda text: "S3BatchUploadError" in text, seconds=30)
|
||||
for index in range(24):
|
||||
_send(owned.gateway, model, key, f"{marker}-{index}")
|
||||
eventually(
|
||||
lambda: owned.log.read_text(),
|
||||
lambda text: "after a failed flush, dropped" in text,
|
||||
seconds=30,
|
||||
)
|
||||
payloads: Final = collect_payloads(sink, 4, seconds=90)
|
||||
landed: Final = frozenset(payload["id"] for payload in payloads)
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == 25
|
||||
assert len(landed) == 4, f"{len(landed)} objects landed with s3_max_queue_size=4"
|
||||
assert f"{marker}-probe" not in landed and f"{marker}-0" not in landed, (
|
||||
f"the oldest events survived the cap: {landed}"
|
||||
)
|
||||
assert f"{marker}-23" in landed, f"the newest event was dropped: {landed}"
|
||||
|
||||
|
||||
def test_s3_v2_retry_age_zero_keeps_aged_object_queued(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3agezero" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RejectingS3Sink(reject_marker=f"{marker}-doomed", reject_status=503, reject_code="SlowDown")
|
||||
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": False, "s3_max_retry_age_seconds": 0})
|
||||
with (
|
||||
owned_proxy_process(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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])
|
||||
_send(owned.gateway, model, key, f"{marker}-doomed")
|
||||
_send_and_wait_until_landed(owned.gateway, model, key, sink, f"{marker}-sibling")
|
||||
_send_and_wait_until_landed(owned.gateway, model, key, sink, f"{marker}-second-flush")
|
||||
_send_and_wait_until_landed(owned.gateway, model, key, sink, f"{marker}-third-flush")
|
||||
attempts_before_clear: Final = sink.rejected_attempts
|
||||
log_text: Final = owned.log.read_text()
|
||||
assert "uploads dropped" not in log_text, log_text
|
||||
assert "retrying longer than" not in log_text, log_text
|
||||
sink.reject_until = time.time()
|
||||
eventually(sink.landed_ids, lambda landed: f"{marker}-doomed" in landed, seconds=60)
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == 4
|
||||
assert attempts_before_clear >= 3, (
|
||||
f"only {attempts_before_clear} PUTs for an object that stayed queued through three delivered flushes; "
|
||||
"with s3_max_retry_age_seconds=0 it must keep retrying longer than any enabled budget"
|
||||
)
|
||||
|
||||
|
||||
def test_s3_v2_throttled_429_object_is_put_once_per_flush(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3throttle" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RejectingS3Sink(reject_marker=f"{marker}-throttled", reject_status=429, reject_code="TooManyRequests")
|
||||
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": False})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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])
|
||||
_send(candidate, model, key, f"{marker}-throttled")
|
||||
_send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-sibling")
|
||||
_send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-second-flush")
|
||||
_send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-third-flush")
|
||||
sink.reject_until = time.time()
|
||||
eventually(sink.landed_ids, lambda landed: f"{marker}-throttled" in landed, seconds=60)
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == 4
|
||||
times: Final = tuple(sink.rejected_times)
|
||||
gaps: Final = tuple(round(later - earlier, 3) for earlier, later in zip(times, times[1:]))
|
||||
assert len(times) >= 3 and min(gaps) >= 1.5, (
|
||||
f"PUTs for a 429 object ran {gaps} apart; the 2 s flush interval allows exactly one attempt per flush "
|
||||
"because 429 is not an in-call retry status"
|
||||
)
|
||||
|
||||
|
||||
def test_s3_v2_default_config_retries_access_denied_and_every_event_lands(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3denied" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RecordingS3Sink(fail_status=403, fail_code="AccessDenied", delay_seconds=0.05)
|
||||
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": False})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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() + 20
|
||||
ids: Final = _push(candidate, model, key, marker, 8)
|
||||
payloads: Final = collect_payloads(sink, 8, seconds=120)
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == 8
|
||||
assert frozenset(payload["id"] for payload in payloads) == ids, (
|
||||
f"a default-config run lost events through a 20s AccessDenied outage: {len(payloads)} landed"
|
||||
)
|
||||
attempt_totals: Final = tuple(sorted(sink.attempt_counts.values()))
|
||||
assert len(attempt_totals) == 8 and all(count >= 4 for count in attempt_totals), (
|
||||
f"each object must see at least one full 3-PUT retry burst before landing: {attempt_totals}"
|
||||
)
|
||||
|
||||
|
||||
@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]
|
||||
|
|
@ -628,3 +907,187 @@ def test_s3_v2_sigterm_mid_burst_loses_only_inflight_without_duplicates(gateway:
|
|||
)
|
||||
targets: Final = tuple(sink.objects())
|
||||
assert len(set(targets)) == len(targets), "the same object was PUT more than once"
|
||||
|
||||
|
||||
RAMP_REQUESTS: Final = 400
|
||||
RAMP_PUT_DELAY_SECONDS: Final = 1.0
|
||||
|
||||
|
||||
def _push(candidate: Gateway, model: str, key: str, marker: str, count: int) -> frozenset[str]:
|
||||
ids: Final = tuple(f"{marker}-{index}" for index in range(count))
|
||||
|
||||
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=64) as pool:
|
||||
returned: Final = frozenset(pool.map(request, ids))
|
||||
assert returned == frozenset(ids)
|
||||
return returned
|
||||
|
||||
|
||||
def test_s3_v2_slow_sink_ramps_concurrency_and_drains_the_backlog(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3ramp" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RecordingS3Sink(delay_seconds=RAMP_PUT_DELAY_SECONDS)
|
||||
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": False, "s3_adaptive_concurrency": True}
|
||||
)
|
||||
with (
|
||||
owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "1", "DEFAULT_S3_BATCH_SIZE": "5000"},
|
||||
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 = _push(candidate, model, key, marker, RAMP_REQUESTS)
|
||||
drain_started: Final = time.monotonic()
|
||||
payloads: Final = collect_payloads(sink, RAMP_REQUESTS, seconds=180)
|
||||
drained_seconds: Final = time.monotonic() - drain_started
|
||||
fixed_sixteen_estimate: Final = RAMP_REQUESTS * RAMP_PUT_DELAY_SECONDS / 16
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == RAMP_REQUESTS
|
||||
assert frozenset(payload["id"] for payload in payloads) == ids
|
||||
assert sink.peak > 16, f"adaptive limiter never ramped past the old fixed bound: peak {sink.peak}"
|
||||
assert drained_seconds < 2 * fixed_sixteen_estimate, (
|
||||
f"backlog of {RAMP_REQUESTS} drained in {drained_seconds:.1f}s with peak concurrency {sink.peak}; "
|
||||
f"even a fixed bound of 16 would need only ~{fixed_sixteen_estimate:.1f}s, so the uploads stalled"
|
||||
)
|
||||
|
||||
|
||||
def test_s3_v2_throttled_sink_halves_in_flight_puts(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3throt" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RecordingS3Sink(fail_status=503, fail_code="SlowDown", delay_seconds=0.3)
|
||||
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": False, "s3_adaptive_concurrency": True, "s3_max_concurrent_uploads": 4},
|
||||
)
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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])
|
||||
healthy_ids: Final = _push(candidate, model, key, f"{marker}-healthy", REQUESTS)
|
||||
collect_payloads(sink, REQUESTS)
|
||||
healthy_peak: Final = sink.peak
|
||||
window_start: Final = time.time()
|
||||
sink.fail_until = window_start + 60
|
||||
throttled_ids: Final = _push(candidate, model, key, f"{marker}-throttled", REQUESTS)
|
||||
first_fail_at: Final = eventually(
|
||||
lambda: next((when for when, _ in sink.attempt_log if when >= window_start), None),
|
||||
lambda when: when is not None,
|
||||
seconds=30,
|
||||
)
|
||||
window_end: Final = first_fail_at + 8.0
|
||||
sink.fail_until = window_end
|
||||
payloads: Final = collect_payloads(sink, 2 * REQUESTS, seconds=120)
|
||||
throttled_peak: Final = sink.peak_between(first_fail_at + 5.0, window_end)
|
||||
throttled_attempts: Final = sum(
|
||||
1 for when, _ in sink.attempt_log if first_fail_at + 5.0 <= when < window_end
|
||||
)
|
||||
assert sum(1 for r in provider.drain() if r.method == "POST") == 2 * REQUESTS
|
||||
assert healthy_peak > 4, (
|
||||
f"healthy peak {healthy_peak} never rose above the configured width 4; nothing to back off from"
|
||||
)
|
||||
assert throttled_attempts > 0, (
|
||||
"no PUTs observed in the measured SlowDown window; the back-off assertion would be vacuous"
|
||||
)
|
||||
assert throttled_peak < healthy_peak, (
|
||||
f"in-flight PUTs during the SlowDown window peaked at {throttled_peak}, not below the healthy peak "
|
||||
f"{healthy_peak}; the limiter did not back off"
|
||||
)
|
||||
assert frozenset(payload["id"] for payload in payloads) == healthy_ids | throttled_ids
|
||||
assert len(sink.objects()) == 2 * REQUESTS
|
||||
|
||||
|
||||
def test_s3_v2_coded_403_is_transient_and_every_id_lands_once(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3coded" + uuid.uuid4().hex[:8]
|
||||
sink: Final = RecordingS3Sink(fail_attempts=3, fail_status=403, fail_code="RequestTimeout", delay_seconds=0.1)
|
||||
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": False})
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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 = _push(candidate, model, key, marker, 4)
|
||||
payloads: Final = collect_payloads(sink, 4)
|
||||
assert frozenset(payload["id"] for payload in payloads) == ids
|
||||
assert sink.attempts >= 7, (
|
||||
f"only {sink.attempts} PUT attempts for 4 objects whose first 3 uploads 403 RequestTimeout; "
|
||||
"coded 403s must be retried"
|
||||
)
|
||||
|
||||
|
||||
def test_s3_v2_success_callback_mode_logs_only_successes(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3succ" + 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,
|
||||
{},
|
||||
{"callbacks": [], "success_callback": ["s3_v2"]},
|
||||
)
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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
|
||||
_send(candidate, model, key, marker)
|
||||
payloads: Final = collect_payloads(sink, 1)
|
||||
assert len(payloads) == 1
|
||||
assert payloads[0]["id"] == marker
|
||||
assert payloads[0]["status"] == "success"
|
||||
|
||||
|
||||
def test_s3_v2_failure_callback_mode_logs_only_failures(gateway: Gateway, tmp_path: Path) -> None:
|
||||
marker: Final = "s3failcb" + 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,
|
||||
{},
|
||||
{"callbacks": [], "failure_callback": ["s3_v2"]},
|
||||
)
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, 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])
|
||||
_send(candidate, model, key, marker)
|
||||
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
|
||||
payloads: Final = collect_payloads(sink, 1)
|
||||
assert len(payloads) == 1
|
||||
assert payloads[0]["status"] == "failure"
|
||||
assert payloads[0]["id"] != marker
|
||||
assert isinstance(payloads[0]["litellm_call_id"], str) and payloads[0]["litellm_call_id"]
|
||||
|
|
|
|||
|
|
@ -322,6 +322,7 @@ class TestS3LoggerAuditLogEvent:
|
|||
logger.s3_path = "my-prefix"
|
||||
logger.log_queue = []
|
||||
logger.batch_size = 100
|
||||
logger.max_queue_size = 100
|
||||
|
||||
audit_log = StandardAuditLogPayload(
|
||||
id="audit-123",
|
||||
|
|
@ -355,6 +356,7 @@ class TestS3LoggerAuditLogEvent:
|
|||
logger.s3_path = None
|
||||
logger.log_queue = []
|
||||
logger.batch_size = 100
|
||||
logger.max_queue_size = 100
|
||||
|
||||
audit_log = StandardAuditLogPayload(
|
||||
id="audit-456",
|
||||
|
|
|
|||
179
tests/unit/integrations/test_adaptive_concurrency.py
Normal file
179
tests/unit/integrations/test_adaptive_concurrency.py
Normal file
|
|
@ -0,0 +1,179 @@
|
|||
import asyncio
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.integrations.adaptive_concurrency import AdaptiveConcurrencyLimiter, PutSample
|
||||
|
||||
_real_sleep: Final = asyncio.sleep
|
||||
|
||||
|
||||
def _limiter(initial: int = 4, floor: int = 1, ceiling: int = 16) -> AdaptiveConcurrencyLimiter:
|
||||
return AdaptiveConcurrencyLimiter(initial=initial, floor=floor, ceiling=ceiling)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_limit_grows_after_limit_clean_samples() -> None:
|
||||
limiter: Final = _limiter(initial=4)
|
||||
for _ in range(4):
|
||||
limiter.record(PutSample(throttled=False))
|
||||
assert limiter.limit == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_limit_does_not_grow_before_the_streak_completes() -> None:
|
||||
limiter: Final = _limiter(initial=4)
|
||||
for _ in range(3):
|
||||
limiter.record(PutSample(throttled=False))
|
||||
assert limiter.limit == 4
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_throttled_sample_halves_the_limit() -> None:
|
||||
limiter: Final = _limiter(initial=16)
|
||||
limiter.record(PutSample(throttled=True))
|
||||
assert limiter.limit == 8
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_limit_clamps_at_the_floor() -> None:
|
||||
limiter: Final = _limiter(initial=4, floor=4)
|
||||
limiter.record(PutSample(throttled=True))
|
||||
assert limiter.limit == 4
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_limit_clamps_at_the_ceiling() -> None:
|
||||
limiter: Final = _limiter(initial=15, ceiling=16)
|
||||
for _ in range(1000):
|
||||
limiter.record(PutSample(throttled=False))
|
||||
assert limiter.limit == 16
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_throttled_sample_resets_the_clean_streak() -> None:
|
||||
limiter: Final = _limiter(initial=4, ceiling=32)
|
||||
for _ in range(3):
|
||||
limiter.record(PutSample(throttled=False))
|
||||
limiter.record(PutSample(throttled=True))
|
||||
limiter.record(PutSample(throttled=False))
|
||||
assert limiter.limit == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_growing_the_limit_wakes_a_waiting_acquirer() -> None:
|
||||
limiter: Final = AdaptiveConcurrencyLimiter(initial=1, floor=1, ceiling=4)
|
||||
acquired: Final[list[str]] = [] # mutable-ok: the waiter task appends to it across the await boundary
|
||||
released: Final = asyncio.Event()
|
||||
|
||||
async def hold() -> None:
|
||||
async with limiter:
|
||||
await released.wait()
|
||||
|
||||
holder: Final = asyncio.create_task(hold())
|
||||
|
||||
async def waiter() -> None:
|
||||
async with limiter:
|
||||
acquired.append("waiter")
|
||||
|
||||
pending: Final = asyncio.create_task(waiter())
|
||||
await _real_sleep(0)
|
||||
assert not acquired
|
||||
|
||||
limiter.record(PutSample(throttled=False))
|
||||
await asyncio.wait_for(asyncio.shield(pending), timeout=5)
|
||||
released.set()
|
||||
await asyncio.wait_for(holder, timeout=5)
|
||||
assert tuple(acquired) == ("waiter",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_releasing_a_slot_wakes_exactly_one_waiter() -> None:
|
||||
limiter: Final = AdaptiveConcurrencyLimiter(initial=1, floor=1, ceiling=4)
|
||||
acquired: Final[list[str]] = [] # mutable-ok: the waiter tasks append to it across the await boundary
|
||||
release: Final = asyncio.Event()
|
||||
|
||||
async def hold() -> None:
|
||||
async with limiter:
|
||||
await _real_sleep(0)
|
||||
|
||||
async def waiter(name: str) -> None:
|
||||
async with limiter:
|
||||
acquired.append(name)
|
||||
await release.wait()
|
||||
|
||||
holder: Final = asyncio.create_task(hold())
|
||||
waiters: Final = tuple(asyncio.create_task(waiter(f"w{i}")) for i in range(3))
|
||||
await _real_sleep(0)
|
||||
await asyncio.wait_for(holder, timeout=5)
|
||||
await _real_sleep(0)
|
||||
assert len(acquired) == 1
|
||||
|
||||
for _ in range(3):
|
||||
limiter.record(PutSample(throttled=False))
|
||||
await _real_sleep(0)
|
||||
assert len(acquired) == 3
|
||||
release.set()
|
||||
await asyncio.gather(*waiters)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_double_cancel_during_release_leaves_in_flight_at_zero() -> None:
|
||||
limiter: Final = AdaptiveConcurrencyLimiter(initial=1, floor=1, ceiling=1)
|
||||
entered: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
|
||||
async def hold() -> None:
|
||||
async with limiter:
|
||||
entered.set()
|
||||
await release.wait()
|
||||
|
||||
holder: Final = asyncio.create_task(hold())
|
||||
await entered.wait()
|
||||
|
||||
waiter: Final = asyncio.create_task(hold())
|
||||
await _real_sleep(0)
|
||||
waiter.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await waiter
|
||||
|
||||
release.set()
|
||||
await asyncio.wait_for(holder, timeout=5)
|
||||
holder.cancel()
|
||||
try:
|
||||
await holder
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert limiter._in_flight == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_cancelled_waiter_is_skipped_when_a_slot_frees() -> None:
|
||||
limiter: Final = AdaptiveConcurrencyLimiter(initial=1, floor=1, ceiling=1)
|
||||
acquired: Final[list[str]] = [] # mutable-ok: waiter tasks append across the await boundary
|
||||
first_entered: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
|
||||
async def hold(name: str, entered: asyncio.Event | None = None) -> None:
|
||||
async with limiter:
|
||||
acquired.append(name)
|
||||
if entered is not None:
|
||||
entered.set()
|
||||
await release.wait()
|
||||
|
||||
holder: Final = asyncio.create_task(hold("holder", first_entered))
|
||||
await first_entered.wait()
|
||||
doomed: Final = asyncio.create_task(hold("doomed"))
|
||||
next_waiter: Final = asyncio.create_task(hold("next"))
|
||||
await _real_sleep(0)
|
||||
doomed.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await doomed
|
||||
release.set()
|
||||
await asyncio.wait_for(holder, timeout=5)
|
||||
await asyncio.wait_for(next_waiter, timeout=5)
|
||||
|
||||
assert "doomed" not in acquired
|
||||
assert "next" in acquired
|
||||
assert limiter._in_flight == 0
|
||||
File diff suppressed because it is too large
Load diff
Loading…
Add table
Reference in a new issue