mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
feat(proxy): add fail_closed_rate_limit_enforcement to reject requests with 503 while Redis rate limit counters are unreachable (#43251)
* feat(proxy): add fail_closed_rate_limit_enforcement to reject requests with 503 while Redis rate limit counters are unreachable * fix(proxy): reject fail-closed rate limit checks before logging the in-memory fallback and pin the boot warning in the lifespan * fix(proxy): coerce the fail-closed flag, fail closed on read-only checks, and refund partial cluster increments * fix(proxy): window-guard rate limit refunds and catch the fail-closed rejection by type * fix(proxy): read the compaction rate-limit gate's limiter from the proxy hook registry * fix(proxy): count the pending request in read-only rate-limit checks and keep the compaction gate off the caller's parallel slot The compaction polyfill's summary-model gate, once it ran against the real v3 limiter, showed two behaviors nobody had chosen. The read-only check compared the stored counter with the same `>` the increment path uses, but a read-only check decides a request that has not been counted yet, so a summary model exactly at its rpm limit still went out. The read-only path now adds the pending increment of 1 before comparing; the increment path is unchanged. The gate also passed the key's max_parallel_requests gauge through, and the read-only gauge count includes the caller's own in-flight slot, so a key with max_parallel_requests: 1 never compacted. The gate now drops that gauge from its descriptors, since the summary call runs inside a request the limiter already admitted. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
41070b1363
commit
9540f19e38
8 changed files with 677 additions and 49 deletions
|
|
@ -30,7 +30,11 @@ from litellm.types.llms.anthropic import (
|
|||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor, RateLimitResponse
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
RateLimitDescriptor,
|
||||
RateLimitDescriptorRateLimitObject,
|
||||
RateLimitResponse,
|
||||
)
|
||||
from litellm.router import Router
|
||||
from litellm.types.llms.anthropic import (
|
||||
AllAnthropicPassThroughMessageValues,
|
||||
|
|
@ -149,6 +153,10 @@ class _CreateOrgRateLimitDescriptors(Protocol):
|
|||
) -> "Sequence[RateLimitDescriptor]": ...
|
||||
|
||||
|
||||
class _GetProxyHook(Protocol):
|
||||
def __call__(self, hook: str) -> object: ...
|
||||
|
||||
|
||||
class _ShouldRateLimit(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
|
|
@ -492,6 +500,24 @@ async def _check_summary_model_budget(
|
|||
return True
|
||||
|
||||
|
||||
def _without_parallel_request_gauges(
|
||||
descriptors: "Sequence[RateLimitDescriptor]",
|
||||
) -> "tuple[RateLimitDescriptor, ...]":
|
||||
return tuple(_without_parallel_request_gauge(descriptor) for descriptor in descriptors)
|
||||
|
||||
|
||||
def _without_parallel_request_gauge(descriptor: "RateLimitDescriptor") -> "RateLimitDescriptor":
|
||||
rate_limit: Final = descriptor.get("rate_limit")
|
||||
if rate_limit is None or rate_limit.get("max_parallel_requests") is None:
|
||||
return descriptor
|
||||
windowed_limits: Final[RateLimitDescriptorRateLimitObject] = {
|
||||
"requests_per_unit": rate_limit.get("requests_per_unit"),
|
||||
"tokens_per_unit": rate_limit.get("tokens_per_unit"),
|
||||
"window_size": rate_limit.get("window_size"),
|
||||
}
|
||||
return {**descriptor, "rate_limit": windowed_limits}
|
||||
|
||||
|
||||
async def _check_summary_model_rate_limit(
|
||||
user_api_key_auth: Optional["UserAPIKeyAuth"],
|
||||
summary_model: str,
|
||||
|
|
@ -508,21 +534,28 @@ async def _check_summary_model_rate_limit(
|
|||
``read_only`` mode so no counter is reserved or incremented — the summary
|
||||
call's actual usage is still charged exactly once by the limiter's
|
||||
post-call success hook (via the propagated ``litellm_metadata``).
|
||||
``max_parallel_requests`` gauges are left out of the check: the summary
|
||||
call runs inside the caller's already admitted request, whose own slot
|
||||
would otherwise count against it.
|
||||
|
||||
Returns True (allow) outside the proxy, when the active limiter does not
|
||||
expose the read-only descriptor check (legacy limiter), or when the
|
||||
descriptor set cannot be built — the only deny signal is a definitive
|
||||
``OVER_LIMIT`` response, so an internal error here forwards the request
|
||||
uncompacted rather than blocking every summary.
|
||||
descriptor set cannot be built — the deny signals are a definitive
|
||||
``OVER_LIMIT`` response and the limiter's own fail-closed rejection
|
||||
(``RateLimitUnverifiableError``, raised when ``fail_closed_rate_limit_enforcement``
|
||||
is on and the counters could not be verified), so any other internal error here
|
||||
forwards the request uncompacted rather than blocking every summary.
|
||||
"""
|
||||
if user_api_key_auth is None:
|
||||
return True
|
||||
try:
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitUnverifiableError
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
except Exception:
|
||||
return True
|
||||
|
||||
limiter: Final[object] = getattr(proxy_logging_obj, "max_parallel_request_limiter", None)
|
||||
get_proxy_hook: Final[_GetProxyHook | None] = getattr(proxy_logging_obj, "get_proxy_hook", None)
|
||||
limiter: Final[object] = get_proxy_hook("parallel_request_limiter") if get_proxy_hook is not None else None
|
||||
should_rate_limit_check: Final[_ShouldRateLimit | None] = getattr(limiter, "should_rate_limit", None)
|
||||
create_descriptors: Final[_CreateRateLimitDescriptors | None] = getattr(
|
||||
limiter, "_create_rate_limit_descriptors", None
|
||||
|
|
@ -566,7 +599,9 @@ async def _check_summary_model_rate_limit(
|
|||
requested_model=summary_model,
|
||||
descriptors=base_descriptors,
|
||||
)
|
||||
descriptors: Final = (*base_descriptors, *create_org_descriptors(user_api_key_auth, summary_model))
|
||||
descriptors: Final = _without_parallel_request_gauges(
|
||||
(*base_descriptors, *create_org_descriptors(user_api_key_auth, summary_model))
|
||||
)
|
||||
if not descriptors:
|
||||
return True
|
||||
parent_otel_span: Final[object] = getattr(user_api_key_auth, "parent_otel_span", None)
|
||||
|
|
@ -575,6 +610,13 @@ async def _check_summary_model_rate_limit(
|
|||
parent_otel_span=parent_otel_span,
|
||||
read_only=True,
|
||||
)
|
||||
except RateLimitUnverifiableError as e:
|
||||
verbose_logger.warning(
|
||||
"compact_20260112: rate-limit counters for summary_model=%s could not be verified; denying: %s",
|
||||
summary_model,
|
||||
e.detail,
|
||||
)
|
||||
return False
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
"compact_20260112: unexpected error during rate-limit check for summary_model=%s; allowing: %s",
|
||||
|
|
|
|||
|
|
@ -2664,6 +2664,14 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
"borrowing the `cache_params` Redis and over the REDIS_* env fallback"
|
||||
),
|
||||
)
|
||||
fail_closed_rate_limit_enforcement: bool | None = Field(
|
||||
None,
|
||||
description=(
|
||||
"reject requests with a 503 while the rate limit counters in Redis are unreachable, instead of "
|
||||
"enforcing tpm/rpm/max_parallel_requests limits per pod from memory (which admits up to N times "
|
||||
"the limit across N pods)"
|
||||
),
|
||||
)
|
||||
control_plane_url: str | None = Field(
|
||||
None,
|
||||
description=(
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ This is currently in development and not yet ready for production.
|
|||
|
||||
import asyncio
|
||||
import binascii
|
||||
import itertools
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
|
|
@ -25,7 +26,9 @@ from typing import (
|
|||
TypedDict,
|
||||
)
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from starlette.status import HTTP_503_SERVICE_UNAVAILABLE
|
||||
from typing_extensions import NotRequired, ReadOnly
|
||||
|
||||
from litellm import DualCache
|
||||
|
|
@ -112,6 +115,44 @@ def _resolve_model_group_alias_via_proxy_router(model: str) -> str | None:
|
|||
return resolve_model_group_alias(llm_router.model_group_alias, model)
|
||||
|
||||
|
||||
FAIL_CLOSED_RATE_LIMIT_ENFORCEMENT_SETTING: Final = "fail_closed_rate_limit_enforcement"
|
||||
RATE_LIMIT_UNVERIFIABLE_MESSAGE: Final = (
|
||||
"Rate limit enforcement unavailable: request counters could not be verified against Redis, and "
|
||||
"fail_closed_rate_limit_enforcement is enabled, so the request was rejected to avoid exceeding the "
|
||||
"configured rate limit. Retry shortly."
|
||||
)
|
||||
|
||||
|
||||
class RateLimitUnverifiableError(HTTPException):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(
|
||||
status_code=HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail={"error": RATE_LIMIT_UNVERIFIABLE_MESSAGE},
|
||||
)
|
||||
|
||||
|
||||
_FAIL_CLOSED_RATE_LIMIT_ENFORCEMENT_FLAG: Final = TypeAdapter(bool | None)
|
||||
|
||||
|
||||
def fail_closed_rate_limit_enforcement_enabled(general_settings: Mapping[str, object]) -> bool:
|
||||
raw_value: Final = general_settings.get(FAIL_CLOSED_RATE_LIMIT_ENFORCEMENT_SETTING)
|
||||
try:
|
||||
return _FAIL_CLOSED_RATE_LIMIT_ENFORCEMENT_FLAG.validate_python(raw_value) is True
|
||||
except ValidationError:
|
||||
verbose_proxy_logger.warning(
|
||||
"general_settings.%s=%r is not a boolean, treating it as disabled",
|
||||
FAIL_CLOSED_RATE_LIMIT_ENFORCEMENT_SETTING,
|
||||
raw_value,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _fail_closed_rate_limit_enforcement_from_general_settings() -> bool:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
return fail_closed_rate_limit_enforcement_enabled(general_settings)
|
||||
|
||||
|
||||
def _sibling_counter_keys(window_key: str) -> tuple[str, str]:
|
||||
prefix: Final = window_key.removesuffix(":window")
|
||||
return f"{prefix}:requests", f"{prefix}:tokens"
|
||||
|
|
@ -156,6 +197,8 @@ end
|
|||
return results
|
||||
"""
|
||||
|
||||
BATCH_COUNTER_READ_SCRIPT: Final = "return redis.call('MGET', unpack(KEYS))"
|
||||
|
||||
CHECK_AND_INCREMENT_BY_N_SCRIPT: Final = """
|
||||
-- Atomic check-and-increment-by-N across one or more descriptors.
|
||||
-- All-or-nothing: if any descriptor would exceed its limit, no counter is
|
||||
|
|
@ -587,6 +630,14 @@ class RequestRateLimiterStash:
|
|||
tpm_limited_tags: frozenset[str] = field(default_factory=frozenset)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CounterRefund:
|
||||
window_key: str
|
||||
counter_key: str
|
||||
window_start: str
|
||||
increment: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TagRateLimit:
|
||||
rpm_limit: int | None
|
||||
|
|
@ -679,6 +730,7 @@ def _parse_output_cap_value(raw_value: object) -> int | None:
|
|||
|
||||
class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
||||
batch_rate_limiter_script: _AsyncLuaScript | None
|
||||
batch_counter_read_script: _AsyncLuaScript | None
|
||||
token_increment_script: _AsyncLuaScript | None
|
||||
check_and_increment_by_n_script: _AsyncLuaScript | None
|
||||
window_guarded_token_increment_script: _AsyncLuaScript | None
|
||||
|
|
@ -692,15 +744,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
time_provider: Callable[[], datetime] | None = None,
|
||||
tag_rate_limit_resolver: TagRateLimitResolver = resolve_tag_rate_limits_from_db,
|
||||
model_group_resolver: Callable[[str], str | None] = _resolve_model_group_alias_via_proxy_router,
|
||||
fail_closed_resolver: Callable[[], bool] = _fail_closed_rate_limit_enforcement_from_general_settings,
|
||||
):
|
||||
self.internal_usage_cache = internal_usage_cache
|
||||
self._time_provider = time_provider or datetime.now
|
||||
self._tag_rate_limit_resolver = tag_rate_limit_resolver
|
||||
self._model_group_resolver = model_group_resolver
|
||||
self._fail_closed_resolver = fail_closed_resolver
|
||||
if self.internal_usage_cache.dual_cache.redis_cache is not None:
|
||||
self.batch_rate_limiter_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
|
||||
BATCH_RATE_LIMITER_SCRIPT
|
||||
)
|
||||
self.batch_counter_read_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
|
||||
BATCH_COUNTER_READ_SCRIPT
|
||||
)
|
||||
self.token_increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
|
||||
TOKEN_INCREMENT_SCRIPT
|
||||
)
|
||||
|
|
@ -723,6 +780,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
else:
|
||||
self.batch_rate_limiter_script = None
|
||||
self.batch_counter_read_script = None
|
||||
self.token_increment_script = None
|
||||
self.check_and_increment_by_n_script = None
|
||||
self.window_guarded_token_increment_script = None
|
||||
|
|
@ -1188,10 +1246,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
keys_to_fetch: list[str],
|
||||
cache_values: CacheCounterValues,
|
||||
key_metadata: dict[str, WindowKeyMetadata],
|
||||
read_only: bool = False,
|
||||
) -> RateLimitResponse:
|
||||
"""
|
||||
Check if the cache values are over the limit.
|
||||
"""
|
||||
pending_increment: Final = 1 if read_only else 0
|
||||
statuses: Final[list[RateLimitStatus]] = []
|
||||
overall_code = "OK"
|
||||
|
||||
|
|
@ -1216,7 +1276,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
if current_limit is None or rate_limit_type is None:
|
||||
continue
|
||||
|
||||
if counter_value is not None and int(counter_value) > current_limit:
|
||||
if counter_value is not None and int(counter_value) + pending_increment > current_limit:
|
||||
overall_code = "OVER_LIMIT"
|
||||
item_code = "OVER_LIMIT"
|
||||
|
||||
|
|
@ -1312,6 +1372,50 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
local_only=True,
|
||||
)
|
||||
|
||||
async def _read_counter_values_from_redis(self, keys: list[str]) -> CacheCounterValues:
|
||||
read_script: Final = self.batch_counter_read_script
|
||||
if read_script is None:
|
||||
return []
|
||||
key_groups: Final = self._group_keys_by_hash_tag(keys)
|
||||
group_values: Final[Sequence[CacheCounterValues]] = [
|
||||
await read_script(keys=group_keys, args=[]) for group_keys in key_groups.values()
|
||||
]
|
||||
values_by_key: Final = dict(
|
||||
zip(
|
||||
itertools.chain.from_iterable(key_groups.values()),
|
||||
itertools.chain.from_iterable(group_values),
|
||||
)
|
||||
)
|
||||
return [values_by_key.get(key) for key in keys]
|
||||
|
||||
async def _read_counter_values_without_incrementing(
|
||||
self,
|
||||
keys: list[str],
|
||||
parent_otel_span: Span | None,
|
||||
) -> CacheCounterValues | None:
|
||||
if self.batch_counter_read_script is None:
|
||||
return await self._batch_get_counter_values(keys=keys, parent_otel_span=parent_otel_span, local_only=False)
|
||||
try:
|
||||
return await self._read_counter_values_from_redis(keys)
|
||||
except Exception as e: # noqa: BLE001 # any Redis/Lua failure degrades to the local mirror unless fail-closed rejects
|
||||
self._reject_if_rate_limit_unverifiable("batch_counter_read_script", e)
|
||||
log_redis_failure(
|
||||
verbose_proxy_logger, logging.WARNING, "batch_counter_read_script failed, using local mirror", e
|
||||
)
|
||||
return await self._batch_get_counter_values(keys=keys, parent_otel_span=parent_otel_span, local_only=True)
|
||||
|
||||
def _reject_if_rate_limit_unverifiable(self, failed_operation: str, error: Exception) -> None:
|
||||
if not self._fail_closed_resolver():
|
||||
return
|
||||
log_redis_failure(
|
||||
verbose_proxy_logger,
|
||||
logging.WARNING,
|
||||
f"fail_closed_rate_limit_enforcement: rejecting request, {failed_operation} could not verify the "
|
||||
"counters against Redis",
|
||||
error,
|
||||
)
|
||||
raise RateLimitUnverifiableError()
|
||||
|
||||
async def _execute_redis_batch_rate_limiter_script(
|
||||
self,
|
||||
keys_to_fetch: list[str],
|
||||
|
|
@ -1330,10 +1434,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
if self.batch_rate_limiter_script is None:
|
||||
return []
|
||||
|
||||
key_groups: Final = self._group_keys_by_hash_tag(keys_to_fetch)
|
||||
key_groups: Final = list(self._group_keys_by_hash_tag(keys_to_fetch).items())
|
||||
all_cache_values: Final[list[CacheCounterValue | None]] = []
|
||||
|
||||
for hash_tag, group_keys in key_groups.items():
|
||||
for index, (hash_tag, group_keys) in enumerate(key_groups):
|
||||
try:
|
||||
group_cache_values: CacheCounterValues = await self.batch_rate_limiter_script(
|
||||
keys=group_keys,
|
||||
|
|
@ -1341,6 +1445,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
all_cache_values.extend(group_cache_values)
|
||||
except Exception as e:
|
||||
if self._fail_closed_resolver():
|
||||
applied_keys = tuple(itertools.chain.from_iterable(keys for _tag, keys in key_groups[:index]))
|
||||
await self._refund_counter_increments(
|
||||
self._counter_refunds_from_batch_values(applied_keys, all_cache_values)
|
||||
)
|
||||
self._reject_if_rate_limit_unverifiable("batch_rate_limiter_script", e)
|
||||
log_redis_failure(
|
||||
verbose_proxy_logger, logging.WARNING, f"Redis Lua script failed for hash tag {hash_tag}", e
|
||||
)
|
||||
|
|
@ -1408,17 +1518,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
|
||||
if cache_values is not None:
|
||||
rate_limit_response: Final = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata)
|
||||
rate_limit_response: Final = self.is_cache_list_over_limit(
|
||||
keys_to_fetch, cache_values, key_metadata, read_only=read_only
|
||||
)
|
||||
if rate_limit_response["overall_code"] == "OVER_LIMIT":
|
||||
return rate_limit_response
|
||||
|
||||
## IF under limit in-memory, check Redis
|
||||
if read_only:
|
||||
# READ-ONLY MODE: Just read current values without incrementing
|
||||
cache_values = await self._batch_get_counter_values( # rebind-ok: read-only mode replaces the in-memory snapshot with Redis values
|
||||
cache_values = await self._read_counter_values_without_incrementing( # rebind-ok: read-only mode replaces the in-memory snapshot with Redis values
|
||||
keys=keys_to_fetch,
|
||||
parent_otel_span=parent_otel_span,
|
||||
local_only=False, # Check Redis too
|
||||
)
|
||||
|
||||
# For keys that don't exist yet, set them to 0
|
||||
|
|
@ -1462,7 +1573,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
window_size=self.window_size,
|
||||
)
|
||||
|
||||
windowed_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata)
|
||||
windowed_response = self.is_cache_list_over_limit(
|
||||
keys_to_fetch, cache_values, key_metadata, read_only=read_only
|
||||
)
|
||||
if windowed_response["overall_code"] == "OVER_LIMIT":
|
||||
return windowed_response
|
||||
|
||||
|
|
@ -1590,7 +1703,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
args=[PARALLEL_REQUEST_SLOT_TTL_SECONDS for _ in gauges],
|
||||
)
|
||||
counts = [max(0, int(value)) for value in raw_counts]
|
||||
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror, never a 500
|
||||
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror unless fail-closed rejects
|
||||
self._reject_if_rate_limit_unverifiable("parallel_count_script", e)
|
||||
log_redis_failure(
|
||||
verbose_proxy_logger, logging.WARNING, "parallel_count_script failed, using local mirror", e
|
||||
)
|
||||
|
|
@ -1623,6 +1737,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
],
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to in-memory enforcement, never a 500
|
||||
self._reject_if_rate_limit_unverifiable("parallel_acquire_script", e)
|
||||
log_redis_failure(
|
||||
verbose_proxy_logger,
|
||||
logging.WARNING,
|
||||
|
|
@ -1941,7 +2056,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
overall_code="OK",
|
||||
statuses=[], # mutable-ok: response contract requires a status list
|
||||
)
|
||||
applied: Final[list[list[AtomicCounterMeta]]] = []
|
||||
applied: Final[list[tuple[CounterRefund, ...]]] = []
|
||||
statuses: Final[list[RateLimitStatus]] = []
|
||||
reservation_windows: Final[set[ReservationWindowIdentity]] = set() # mutable-ok: filled by the group loop
|
||||
raw: list[CacheCounterValue]
|
||||
|
|
@ -1957,15 +2072,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
# state ambiguous. Refund any prior groups so Redis returns
|
||||
# to its pre-call state, then fall back to in-memory for the
|
||||
# whole call (counters there are independent of Redis).
|
||||
await self._refund_applied_descriptor_groups(applied)
|
||||
self._reject_if_rate_limit_unverifiable("check_and_increment_by_n_script", e)
|
||||
log_redis_failure(
|
||||
verbose_proxy_logger,
|
||||
logging.ERROR,
|
||||
f"atomic_check_and_increment_by_n: Redis Lua execution failed ({type(e).__name__}). Refunding "
|
||||
f"atomic_check_and_increment_by_n: Redis Lua execution failed ({type(e).__name__}). Refunded "
|
||||
f"{len(applied)} prior descriptors and falling back to in-memory enforcement, counters will "
|
||||
f"diverge from Redis until window expires (window_size={self.window_size}s)",
|
||||
e,
|
||||
)
|
||||
await self._refund_applied_descriptor_groups(applied)
|
||||
flat_meta: list[AtomicCounterMeta] = [m for _k, _a, group_meta in descriptor_groups for m in group_meta]
|
||||
async with self._check_and_increment_lock:
|
||||
return await self._atomic_check_and_increment_in_memory(
|
||||
|
|
@ -1979,7 +2095,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
return response
|
||||
if len(descriptor_groups) == 1:
|
||||
return response
|
||||
applied.append(meta)
|
||||
applied.append(self._counter_refunds_from_atomic_response(raw, meta))
|
||||
statuses.extend(response["statuses"])
|
||||
reservation_windows.update(response.get("reservation_windows", frozenset()))
|
||||
|
||||
|
|
@ -1991,32 +2107,63 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
|
||||
async def _refund_applied_descriptor_groups(
|
||||
self,
|
||||
applied: list[list[AtomicCounterMeta]],
|
||||
applied: Sequence[Sequence[CounterRefund]],
|
||||
) -> None:
|
||||
"""
|
||||
Decrement counters for descriptor groups already applied via Lua.
|
||||
Best-effort: refund failures are logged but not raised — the original
|
||||
OVER_LIMIT / fallback decision is what matters to the caller.
|
||||
"""
|
||||
if not applied:
|
||||
await self._refund_counter_increments(tuple(itertools.chain.from_iterable(applied)))
|
||||
|
||||
@staticmethod
|
||||
def _counter_refunds_from_atomic_response(
|
||||
raw: Sequence[CacheCounterValue],
|
||||
per_counter_meta: Sequence[AtomicCounterMeta],
|
||||
) -> tuple[CounterRefund, ...]:
|
||||
return tuple(
|
||||
CounterRefund(
|
||||
window_key=meta["window_key"],
|
||||
counter_key=meta["counter_key"],
|
||||
window_start=str(int(raw[2 + index * 2])),
|
||||
increment=meta["increment"],
|
||||
)
|
||||
for index, meta in enumerate(per_counter_meta)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _counter_refunds_from_batch_values(
|
||||
applied_keys: Sequence[str],
|
||||
applied_values: Sequence[CacheCounterValue | None],
|
||||
) -> tuple[CounterRefund, ...]:
|
||||
pairs: Final = tuple(zip(range(0, len(applied_keys), 2), applied_values[::2]))
|
||||
return tuple(
|
||||
CounterRefund(
|
||||
window_key=applied_keys[offset],
|
||||
counter_key=applied_keys[offset + 1],
|
||||
window_start=str(int(window_start)),
|
||||
increment=1,
|
||||
)
|
||||
for offset, window_start in pairs
|
||||
if window_start is not None
|
||||
)
|
||||
|
||||
async def _refund_counter_increments(self, refunds: Sequence[CounterRefund]) -> None:
|
||||
if self.window_guarded_token_increment_script is None:
|
||||
return
|
||||
redis_cache: Final = self.internal_usage_cache.dual_cache.redis_cache
|
||||
if redis_cache is None:
|
||||
return
|
||||
for group_meta in applied:
|
||||
for entry in group_meta:
|
||||
try:
|
||||
await redis_cache.async_increment(
|
||||
key=entry["counter_key"],
|
||||
value=-entry["increment"],
|
||||
)
|
||||
except Exception as e:
|
||||
log_redis_failure(
|
||||
verbose_proxy_logger,
|
||||
logging.WARNING,
|
||||
f"Failed to refund {entry['counter_key']} on cross-descriptor rollback",
|
||||
e,
|
||||
)
|
||||
for refund in refunds:
|
||||
try:
|
||||
await self.window_guarded_token_increment_script(
|
||||
keys=[refund.window_key, refund.counter_key], # mutable-ok: Redis script API takes a list
|
||||
args=[refund.window_start, -refund.increment, 0], # mutable-ok: Redis script API takes a list
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best-effort rollback, the rejection already decided the request
|
||||
log_redis_failure(
|
||||
verbose_proxy_logger,
|
||||
logging.WARNING,
|
||||
f"Failed to refund {refund.counter_key} on rollback",
|
||||
e,
|
||||
)
|
||||
|
||||
def _build_atomic_response(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -545,6 +545,7 @@ from litellm.proxy.health_endpoints._health_endpoints import router as health_ro
|
|||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||
_PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
)
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import fail_closed_rate_limit_enforcement_enabled
|
||||
from litellm.proxy.hooks.prompt_injection_detection import (
|
||||
_OPTIONAL_PromptInjectionDetection,
|
||||
)
|
||||
|
|
@ -1471,6 +1472,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
|
|||
max_budget=litellm.max_budget,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
ProxyStartupEvent._warn_fail_closed_rate_limits_without_redis(
|
||||
fail_closed_rate_limit_enforcement=fail_closed_rate_limit_enforcement_enabled(general_settings),
|
||||
redis_usage_cache=redis_usage_cache,
|
||||
)
|
||||
|
||||
### START BATCH WRITING DB + CHECKING NEW MODELS###
|
||||
worker_heartbeat: Final = (
|
||||
|
|
@ -9827,6 +9832,20 @@ class ProxyStartupEvent:
|
|||
max_budget,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _warn_fail_closed_rate_limits_without_redis(
|
||||
fail_closed_rate_limit_enforcement: bool, redis_usage_cache: RedisCache | None
|
||||
) -> None:
|
||||
if redis_usage_cache is not None or not fail_closed_rate_limit_enforcement:
|
||||
return
|
||||
|
||||
verbose_proxy_logger.warning(
|
||||
"general_settings.fail_closed_rate_limit_enforcement is enabled but no Redis is configured, so rate "
|
||||
"limits are enforced per pod from memory and the setting rejects nothing. Configure "
|
||||
"general_settings.coordination_redis (or REDIS_HOST/REDIS_PORT/REDIS_PASSWORD) to share the counters "
|
||||
"across pods and make the setting effective."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _initialize_startup_logging(
|
||||
cls,
|
||||
|
|
|
|||
|
|
@ -6718,6 +6718,262 @@ async def test_an_open_circuit_breaker_reads_the_sliding_window_locally_without_
|
|||
assert any("circuit breaker is open" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
class _UnreachableRedis:
|
||||
def async_register_script(self, script: str):
|
||||
async def refused(keys, args):
|
||||
raise ConnectionError("Error 61 connecting to 127.0.0.1:6379. Connection refused.")
|
||||
|
||||
return refused
|
||||
|
||||
|
||||
class _ScriptedRedis:
|
||||
def __init__(
|
||||
self,
|
||||
failing_script: str | None = None,
|
||||
failing_batch_call: int | None = None,
|
||||
stored_counter_value: int = 0,
|
||||
):
|
||||
self.failing_script = failing_script
|
||||
self.failing_batch_call = failing_batch_call
|
||||
self.stored_counter_value = stored_counter_value
|
||||
self.released_slots: list[tuple[list[str], list[str]]] = []
|
||||
self.batch_calls = 0
|
||||
self.batch_call_keys: list[list[str]] = []
|
||||
self.batch_call_args: list[list[object]] = []
|
||||
self.increments: list[tuple[str, float]] = []
|
||||
self.guarded_increments: list[tuple[list[str], list[object]]] = []
|
||||
|
||||
async def async_increment(self, key: str, value: float, **kwargs):
|
||||
self.increments.append((key, value))
|
||||
return value
|
||||
|
||||
def async_register_script(self, script: str):
|
||||
from litellm.proxy.hooks import parallel_request_limiter_v3 as v3
|
||||
|
||||
async def run(keys, args):
|
||||
if script == self.failing_script:
|
||||
raise ConnectionError("Error 61 connecting to 127.0.0.1:6379. Connection refused.")
|
||||
if script == v3.WINDOW_GUARDED_TOKEN_INCREMENT_SCRIPT:
|
||||
self.guarded_increments.append((list(keys), list(args)))
|
||||
return [1, 0] * (len(keys) // 2)
|
||||
if script == v3.BATCH_RATE_LIMITER_SCRIPT:
|
||||
self.batch_calls += 1
|
||||
self.batch_call_keys.append(list(keys))
|
||||
self.batch_call_args.append(list(args))
|
||||
if self.batch_calls == self.failing_batch_call:
|
||||
raise ConnectionError("Error 61 connecting to 127.0.0.1:6379. Connection refused.")
|
||||
return [args[0], self.batch_calls] * (len(keys) // 2)
|
||||
if script == v3.BATCH_COUNTER_READ_SCRIPT:
|
||||
return [int(time.time()) if key.endswith(":window") else self.stored_counter_value for key in keys]
|
||||
if script == v3.PARALLEL_COUNT_SCRIPT:
|
||||
return [0 for _ in keys]
|
||||
if script == v3.PARALLEL_ACQUIRE_SCRIPT:
|
||||
return [0, *[1 for _ in keys]]
|
||||
if script == v3.PARALLEL_RELEASE_SCRIPT:
|
||||
self.released_slots.append((list(keys), list(args)))
|
||||
return [0 for _ in keys]
|
||||
raise AssertionError(f"unexpected script: {script[:60]}")
|
||||
|
||||
return run
|
||||
|
||||
|
||||
def _handler_with_redis(redis, fail_closed: bool | None = None):
|
||||
internal_usage_cache = InternalUsageCache(DualCache(redis_cache=redis)) # pyright: ignore[reportArgumentType] # duck-typed Redis double
|
||||
if fail_closed is None:
|
||||
return _PROXY_MaxParallelRequestsHandler(internal_usage_cache=internal_usage_cache)
|
||||
return _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=internal_usage_cache,
|
||||
fail_closed_resolver=lambda: fail_closed,
|
||||
)
|
||||
|
||||
|
||||
async def _admit(handler, auth, data=None):
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
cache=handler.internal_usage_cache.dual_cache,
|
||||
data=data if data is not None else {"model": "test-model", "messages": [{"role": "user", "content": "hi"}]},
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
|
||||
async def _read_only_check(handler, auth):
|
||||
descriptors = handler._create_rate_limit_descriptors(
|
||||
user_api_key_dict=auth,
|
||||
data={"model": "test-model"},
|
||||
rpm_limit_type=None,
|
||||
tpm_limit_type=None,
|
||||
model_has_failures=False,
|
||||
)
|
||||
return await handler.should_rate_limit(descriptors=descriptors, read_only=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"limits",
|
||||
[{"rpm_limit": 2}, {"max_parallel_requests": 1}, {"tpm_limit": 1000}],
|
||||
ids=["rpm_window", "parallel_gauge", "tpm_reservation"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_closed_rejects_with_503_when_redis_counters_are_unreachable(limits):
|
||||
handler = _handler_with_redis(_UnreachableRedis(), fail_closed=True)
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed"), **limits)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _admit(handler, auth)
|
||||
|
||||
assert exc.value.status_code == 503
|
||||
assert not isinstance(exc.value, ProxyRateLimitError)
|
||||
assert "fail_closed_rate_limit_enforcement" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_open_default_keeps_enforcing_per_pod_from_memory_when_redis_counters_are_unreachable():
|
||||
handler = _handler_with_redis(_UnreachableRedis(), fail_closed=False)
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-open"), rpm_limit=2)
|
||||
|
||||
await _admit(handler, auth)
|
||||
await _admit(handler, auth)
|
||||
with pytest.raises(ProxyRateLimitError) as exc:
|
||||
await _admit(handler, auth)
|
||||
|
||||
assert exc.value.status_code == 429
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_closed_is_a_no_op_while_redis_answers():
|
||||
handler = _handler_with_redis(_ScriptedRedis(), fail_closed=True)
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed-healthy"), rpm_limit=2)
|
||||
|
||||
await _admit(handler, auth)
|
||||
await _admit(handler, auth)
|
||||
with pytest.raises(ProxyRateLimitError) as exc:
|
||||
await _admit(handler, auth)
|
||||
|
||||
assert exc.value.status_code == 429
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_closed_tpm_rejection_releases_the_parallel_slot_it_acquired():
|
||||
from litellm.proxy.hooks import parallel_request_limiter_v3 as v3
|
||||
|
||||
redis = _ScriptedRedis(failing_script=v3.CHECK_AND_INCREMENT_BY_N_SCRIPT)
|
||||
handler = _handler_with_redis(redis, fail_closed=True)
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed-slot"), max_parallel_requests=1, tpm_limit=1000)
|
||||
data = {"model": "test-model", "messages": [{"role": "user", "content": "hi"}]}
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _admit(handler, auth, data)
|
||||
assert exc.value.status_code == 503
|
||||
acquired = get_or_create_request_stash().parallel_slot
|
||||
assert acquired is not None
|
||||
|
||||
await handler.async_post_call_failure_hook(
|
||||
request_data=data, original_exception=exc.value, user_api_key_dict=auth
|
||||
)
|
||||
|
||||
assert redis.released_slots == [(list(acquired["counter_keys"]), [acquired["slot_id"]])]
|
||||
assert get_or_create_request_stash().parallel_slot is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_closed_rate_limit_enforcement_is_read_from_general_settings(monkeypatch):
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed-settings"), rpm_limit=2)
|
||||
|
||||
monkeypatch.setitem(proxy_server.general_settings, "fail_closed_rate_limit_enforcement", True)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _admit(_handler_with_redis(_UnreachableRedis()), auth)
|
||||
assert exc.value.status_code == 503
|
||||
|
||||
monkeypatch.delitem(proxy_server.general_settings, "fail_closed_rate_limit_enforcement")
|
||||
await _admit(_handler_with_redis(_UnreachableRedis()), auth)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"configured_value, rejects",
|
||||
[(True, True), ("true", True), (False, False), ("false", False), ("sometimes", False)],
|
||||
ids=["bool_true", "string_true", "bool_false", "string_false", "not_a_boolean"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_closed_rate_limit_enforcement_coerces_the_general_settings_value(
|
||||
monkeypatch, configured_value, rejects
|
||||
):
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed-coerced"), rpm_limit=2)
|
||||
monkeypatch.setitem(proxy_server.general_settings, "fail_closed_rate_limit_enforcement", configured_value)
|
||||
|
||||
if not rejects:
|
||||
await _admit(_handler_with_redis(_UnreachableRedis()), auth)
|
||||
return
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _admit(_handler_with_redis(_UnreachableRedis()), auth)
|
||||
assert exc.value.status_code == 503
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"limits",
|
||||
[{"rpm_limit": 2}, {"max_parallel_requests": 1}],
|
||||
ids=["rpm_window", "parallel_gauge"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_fail_closed_read_only_check_rejects_with_503_when_redis_counters_are_unreachable(limits):
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed-read-only"), **limits)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _read_only_check(_handler_with_redis(_UnreachableRedis(), fail_closed=True), auth)
|
||||
assert exc.value.status_code == 503
|
||||
|
||||
response = await _read_only_check(_handler_with_redis(_UnreachableRedis(), fail_closed=False), auth)
|
||||
assert response["overall_code"] == "OK"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stored_counter_value, expected_code", [(1, "OK"), (2, "OVER_LIMIT"), (3, "OVER_LIMIT")])
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_only_check_reports_the_redis_counters_without_incrementing_them(
|
||||
stored_counter_value, expected_code
|
||||
):
|
||||
redis = _ScriptedRedis(stored_counter_value=stored_counter_value)
|
||||
auth = UserAPIKeyAuth(api_key=hash_token("sk-read-only-counters"), rpm_limit=2)
|
||||
|
||||
response = await _read_only_check(_handler_with_redis(redis, fail_closed=True), auth)
|
||||
|
||||
assert response["overall_code"] == expected_code
|
||||
assert redis.batch_calls == 0
|
||||
assert redis.increments == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fail_closed", [True, False], ids=["fail_closed", "fail_open"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_increment_refunds_counters_already_applied_when_a_later_cluster_slot_fails(fail_closed):
|
||||
from unittest.mock import patch
|
||||
|
||||
redis = _ScriptedRedis(failing_batch_call=2)
|
||||
handler = _handler_with_redis(redis, fail_closed=fail_closed)
|
||||
auth = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-cluster-partial"), rpm_limit=5, user_id="cluster-user", user_rpm_limit=5
|
||||
)
|
||||
|
||||
with patch.object(handler, "_is_redis_cluster", return_value=True):
|
||||
if fail_closed:
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _admit(handler, auth)
|
||||
assert exc.value.status_code == 503
|
||||
else:
|
||||
await _admit(handler, auth)
|
||||
|
||||
assert len(redis.batch_call_keys) == 2
|
||||
applied_keys = redis.batch_call_keys[0]
|
||||
assert applied_keys
|
||||
window_start_at_increment = str(redis.batch_call_args[0][0])
|
||||
expected_refunds = [
|
||||
([applied_keys[offset], applied_keys[offset + 1]], [window_start_at_increment, -1, 0])
|
||||
for offset in range(0, len(applied_keys), 2)
|
||||
]
|
||||
assert redis.guarded_increments == (expected_refunds if fail_closed else [])
|
||||
assert redis.increments == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"limits, request_data, counter_scope",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -952,6 +952,49 @@ def test_startup_does_not_warn_without_global_budget(caplog, max_budget):
|
|||
assert "litellm.max_budget" not in caplog.text
|
||||
|
||||
|
||||
def test_startup_warns_for_fail_closed_rate_limits_without_redis(caplog):
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
ProxyStartupEvent._warn_fail_closed_rate_limits_without_redis(
|
||||
fail_closed_rate_limit_enforcement=True, redis_usage_cache=None
|
||||
)
|
||||
|
||||
assert "fail_closed_rate_limit_enforcement" in caplog.text
|
||||
assert "rejects nothing" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fail_closed, redis_usage_cache", [(True, MagicMock()), (False, None)])
|
||||
def test_startup_does_not_warn_for_fail_closed_rate_limits_when_nothing_is_lost(caplog, fail_closed, redis_usage_cache):
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
ProxyStartupEvent._warn_fail_closed_rate_limits_without_redis(
|
||||
fail_closed_rate_limit_enforcement=fail_closed, redis_usage_cache=redis_usage_cache
|
||||
)
|
||||
|
||||
assert "fail_closed_rate_limit_enforcement" not in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_startup_event_warns_for_fail_closed_rate_limits_without_redis(caplog):
|
||||
scheduler = AsyncIOScheduler()
|
||||
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} | {
|
||||
"LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY": "true"
|
||||
}
|
||||
with (
|
||||
patch.dict(os.environ, clean_env, clear=True),
|
||||
patch.object(ps, "scheduler", scheduler),
|
||||
patch.dict(ps.general_settings, {"fail_closed_rate_limit_enforcement": True}),
|
||||
caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"),
|
||||
):
|
||||
try:
|
||||
async with proxy_startup_event(app=None):
|
||||
pass
|
||||
finally:
|
||||
if scheduler.running:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
assert "fail_closed_rate_limit_enforcement" in caplog.text
|
||||
assert "rejects nothing" in caplog.text
|
||||
|
||||
|
||||
def test_proxy_startup_event_warns_for_global_budget_without_database():
|
||||
"""Pin the lifespan call that prevents silent DB-less budgets.
|
||||
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.experimental_pass_through.context_management import (
|
||||
|
|
@ -36,6 +37,7 @@ from litellm.llms.anthropic.experimental_pass_through.context_management.editors
|
|||
from litellm.llms.anthropic.experimental_pass_through.context_management.result import (
|
||||
PolyfillResult,
|
||||
)
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitUnverifiableError
|
||||
|
||||
MODEL = "openai/gpt-4o"
|
||||
|
||||
|
|
@ -1765,12 +1767,25 @@ async def test_summary_model_allowed_when_within_model_budget():
|
|||
assert not result.applied_edits[0].get("error")
|
||||
|
||||
|
||||
class _LegacyLimiter:
|
||||
async def async_pre_call_hook(self, **kwargs):
|
||||
return None
|
||||
|
||||
|
||||
def _proxy_logging_like_the_live_proxy(active_limiter: object) -> MagicMock:
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.max_parallel_request_limiter = _LegacyLimiter()
|
||||
proxy_logging.get_proxy_hook = lambda hook: active_limiter if hook == "parallel_request_limiter" else None
|
||||
return proxy_logging
|
||||
|
||||
|
||||
class _FakeRateLimiter:
|
||||
"""Minimal stand-in for ``_PROXY_MaxParallelRequestsHandler_v3`` exposing
|
||||
just the descriptor-build + read-only check surface the editor consults."""
|
||||
|
||||
def __init__(self, overall_code: str):
|
||||
def __init__(self, overall_code: str, raises: Exception | None = None):
|
||||
self._overall_code = overall_code
|
||||
self._raises = raises
|
||||
self.read_only_checked = False
|
||||
|
||||
def _create_rate_limit_descriptors(self, **kwargs):
|
||||
|
|
@ -1793,9 +1808,62 @@ class _FakeRateLimiter:
|
|||
|
||||
async def should_rate_limit(self, **kwargs):
|
||||
self.read_only_checked = kwargs.get("read_only") is True
|
||||
if self._raises is not None:
|
||||
raise self._raises
|
||||
return {"overall_code": self._overall_code}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"limiter_error, summary_called",
|
||||
[
|
||||
(RateLimitUnverifiableError(), False),
|
||||
(HTTPException(status_code=500, detail="unrelated proxy error"), True),
|
||||
(RuntimeError("descriptor build exploded"), True),
|
||||
],
|
||||
ids=["fail_closed_rejection_denies", "other_http_error_allows", "internal_error_allows"],
|
||||
)
|
||||
async def test_summary_model_rate_limit_check_errors(limiter_error, summary_called):
|
||||
"""The limiter's fail-closed 503 is a verdict and skips the summary call the
|
||||
way OVER_LIMIT does; any other error keeps failing open."""
|
||||
messages = _simple_messages()
|
||||
mock_call = AsyncMock(return_value=_make_mock_response("<summary>ok</summary>"))
|
||||
|
||||
auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
|
||||
limiter = _FakeRateLimiter("OK", raises=limiter_error)
|
||||
proxy_logging = _proxy_logging_like_the_live_proxy(limiter)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
|
||||
return_value="claude-haiku-4-5",
|
||||
),
|
||||
patch("litellm.token_counter", return_value=200_000),
|
||||
patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
|
||||
mock_call,
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging),
|
||||
):
|
||||
result = await apply_compact_20260112(
|
||||
model=MODEL,
|
||||
messages=messages,
|
||||
tools=None,
|
||||
system=None,
|
||||
edit_spec=_EDIT_SPEC_DEFAULT,
|
||||
user_api_key_auth=auth,
|
||||
)
|
||||
|
||||
assert limiter.read_only_checked is True
|
||||
if summary_called:
|
||||
mock_call.assert_awaited_once()
|
||||
assert result.compaction_block is not None
|
||||
assert not result.applied_edits[0].get("error")
|
||||
return
|
||||
mock_call.assert_not_awaited()
|
||||
assert result.compaction_block is None
|
||||
assert result.applied_edits[0].get("error") == "summary_model_rate_limit_exceeded"
|
||||
|
||||
|
||||
async def test_summary_model_denied_when_over_rate_limit():
|
||||
"""A caller already at their configured RPM/TPM for the summary model cannot
|
||||
drive an extra summary completion via compaction."""
|
||||
|
|
@ -1804,8 +1872,7 @@ async def test_summary_model_denied_when_over_rate_limit():
|
|||
|
||||
auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
|
||||
limiter = _FakeRateLimiter("OVER_LIMIT")
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.max_parallel_request_limiter = limiter
|
||||
proxy_logging = _proxy_logging_like_the_live_proxy(limiter)
|
||||
|
||||
with (
|
||||
patch(
|
||||
|
|
@ -1841,8 +1908,7 @@ async def test_summary_model_allowed_when_within_rate_limit():
|
|||
|
||||
auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
|
||||
limiter = _FakeRateLimiter("OK")
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.max_parallel_request_limiter = limiter
|
||||
proxy_logging = _proxy_logging_like_the_live_proxy(limiter)
|
||||
|
||||
with (
|
||||
patch(
|
||||
|
|
@ -1871,6 +1937,53 @@ async def test_summary_model_allowed_when_within_rate_limit():
|
|||
assert not result.applied_edits[0].get("error")
|
||||
|
||||
|
||||
async def test_summary_model_allowed_while_the_caller_holds_the_keys_only_parallel_slot():
|
||||
"""The summary call runs inside a request the limiter already admitted, so the
|
||||
caller's own in-flight slot must not trip a ``max_parallel_requests`` gauge."""
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3
|
||||
from litellm.proxy.utils import InternalUsageCache, hash_token
|
||||
|
||||
messages = _simple_messages()
|
||||
mock_call = AsyncMock(return_value=_make_mock_response("<summary>ok</summary>"))
|
||||
limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache()))
|
||||
auth = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-compact-parallel-slot"), max_parallel_requests=1, models=["all-proxy-models"]
|
||||
)
|
||||
await limiter.async_pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
cache=limiter.internal_usage_cache.dual_cache,
|
||||
data={"model": MODEL, "messages": messages},
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
|
||||
return_value="claude-haiku-4-5",
|
||||
),
|
||||
patch("litellm.token_counter", return_value=200_000),
|
||||
patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
|
||||
mock_call,
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", _proxy_logging_like_the_live_proxy(limiter)),
|
||||
):
|
||||
result = await apply_compact_20260112(
|
||||
model=MODEL,
|
||||
messages=messages,
|
||||
tools=None,
|
||||
system=None,
|
||||
edit_spec=_EDIT_SPEC_DEFAULT,
|
||||
user_api_key_auth=auth,
|
||||
)
|
||||
|
||||
mock_call.assert_awaited_once()
|
||||
assert result.compaction_block is not None
|
||||
assert not result.applied_edits[0].get("error")
|
||||
|
||||
|
||||
async def test_summary_model_rate_limit_skipped_for_legacy_limiter():
|
||||
"""A limiter without the v3 read-only check surface fails open so the summary
|
||||
call still proceeds (its usage is still charged post-call)."""
|
||||
|
|
@ -1879,12 +1992,7 @@ async def test_summary_model_rate_limit_skipped_for_legacy_limiter():
|
|||
|
||||
auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
|
||||
|
||||
class _LegacyLimiter:
|
||||
async def async_pre_call_hook(self, **kwargs):
|
||||
return None
|
||||
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.max_parallel_request_limiter = _LegacyLimiter()
|
||||
proxy_logging = _proxy_logging_like_the_live_proxy(_LegacyLimiter())
|
||||
|
||||
with (
|
||||
patch(
|
||||
|
|
|
|||
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -28214,6 +28214,11 @@ export interface components {
|
|||
* @description If True, router fallbacks configured in router_settings are only attempted when the calling key (and its team and project) is allowed to call the fallback model; unauthorized fallback targets are skipped and the primary model's error is returned. Default is False.
|
||||
*/
|
||||
enforce_fallback_model_access?: boolean | null;
|
||||
/**
|
||||
* Fail Closed Rate Limit Enforcement
|
||||
* @description reject requests with a 503 while the rate limit counters in Redis are unreachable, instead of enforcing tpm/rpm/max_parallel_requests limits per pod from memory (which admits up to N times the limit across N pods)
|
||||
*/
|
||||
fail_closed_rate_limit_enforcement?: boolean | null;
|
||||
/**
|
||||
* Failed Login Block Seconds
|
||||
* @description How long a blocked source address, or source address and username, stays blocked. Every attempt from a blocked key, right or wrong, is refused with 429 before the password is checked; the block is not extended by refused attempts. Set under `general_settings` in config.yaml. Defaults to 300
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue