scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

fix(rate-limiting): refresh the global hook's own concurrency key ttl on every admission

Browse source 
veria-ai finding on the previous commit: TAG_RL_CHECK_AND_INCR_SCRIPT is
shared with model_based_tag_rate_limits_hook, but only that hook's own
Redis and in-memory call sites were updated to carry refresh_ttl through.
global_tag_rate_limits_hook's own _check_and_increment_one still called the
script with three args and never passed refresh_ttl to InMemoryCache, so a
global concurrency bucket's ttl stayed fixed from its first admission and
could expire mid-flight under sustained traffic, admitting past the cap.

Also fixes ANN001/reportArgumentType drift the async_log_success_event and
async_log_failure_event annotations surfaced: kwargs/response_obj are now
typed, litellm_params_for_metadata is narrowed via isinstance instead of an
unchecked assumption, and a stray tuple() wrap that never matched its
list-typed parameter is dropped.
This commit is contained in:
Deepanshu 2026-08-27 17:29:54 -04:00
parent ffd00dd4aa
commit 5d0eca67e8
2 changed files with 143 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
litellm/proxy/hooks/global_tag_rate_limits_hook.py
View file

@ -120,7 +120,6 @@ from litellm.router_strategy.tag_based_routing import (
)
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.router import TagRateLimitEntry, TagRateLimits
from litellm.types.utils import StandardLoggingPayload
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -327,10 +326,12 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
return built
async def _check_and_increment_one(
self, cache: InternalUsageCache, key: str, limit: float, increment: float, ttl: int
self, cache: InternalUsageCache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool
) -> tuple[bool, float]:
if self._check_and_incr_script is not None:
raw: Final = await self._check_and_incr_script(keys=(key,), args=(limit, increment, ttl))
raw: Final = await self._check_and_incr_script(
keys=(key,), args=(limit, increment, ttl, 1 if refresh_ttl else 0)
)
return bool(raw[0]), float(raw[1])
async with self._lock:
current_value: Final = await cache.async_get_cache(key=key, litellm_parent_otel_span=None)
@ -338,7 +339,9 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
if current + increment > limit:
return False, current
new_value: Final = current + increment
await cache.async_set_cache(key=key, value=new_value, ttl=ttl, litellm_parent_otel_span=None)
await cache.async_set_cache(
key=key, value=new_value, ttl=ttl, refresh_ttl=refresh_ttl, litellm_parent_otel_span=None
)
return True, new_value
async def _decrement_floor_zero(self, cache: InternalUsageCache, key: str, delta: float) -> None:
@ -352,17 +355,17 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
async def _atomic_check_and_increment(
self,
checks: Sequence[tuple[InternalUsageCache, str, float, float, int]],
checks: Sequence[tuple[InternalUsageCache, str, float, float, int, bool]],
) -> tuple[int | None, tuple[float, ...]]:
"""All-or-nothing atomic admission across `checks`: on a rejection,
refunds every check admitted earlier in this batch."""
if not checks:
return None, ()
admitted_values: Final = [] # mutable-ok: sequential async accumulator, discardable on early rejection
for index, (cache, key, limit, increment, ttl) in enumerate(checks):
for index, (cache, key, limit, increment, ttl, refresh_ttl) in enumerate(checks):
admitted = False
try:
admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl)
admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl)
finally:
if not admitted:
await self._refund_admitted(checks, up_to_index=index)
@ -373,10 +376,10 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
return None, tuple(admitted_values)
async def _refund_admitted(
self, checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], up_to_index: int
self, checks: Sequence[tuple[InternalUsageCache, str, float, float, int, bool]], up_to_index: int
) -> None:
for refund_index in range(up_to_index):
refund_cache, refund_key, _limit, refund_increment, _ttl = checks[refund_index]
refund_cache, refund_key, _limit, refund_increment, _ttl, _refresh_ttl = checks[refund_index]
try:
await self._decrement_floor_zero(refund_cache, refund_key, -refund_increment)
except Exception as e: # noqa: BLE001 - one failed refund must not block refunding the rest
@ -597,6 +600,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
)
else 1.0,
self._ttl_for(check.unit, check.entry),
check.unit == "concurrency",
)
for partition, check in zip(atomic_partitions, atomic_checks)
)
@ -667,7 +671,13 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
"""
await self._release_pending_for_call_id(request_data)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
async def async_log_failure_event(
self,
kwargs: Mapping[str, object],
response_obj: object,
start_time: datetime | None,
end_time: datetime | None,
) -> None:
# Always release regardless of which hook raised: this hook's own
# rejection never reserves a slot, so pending_concurrency_keys is
# already empty in that case and the check below no-ops; a rejection
@ -675,7 +685,13 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
# land after this hook already reserved its own slot.
await self._release_pending_for_call_id(kwargs)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
async def async_log_success_event(
self,
kwargs: Mapping[str, object],
response_obj: object,
start_time: datetime | None,
end_time: datetime | None,
) -> None:
stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs))
if stash is not None and stash.pending_concurrency_keys:
release_keys: Final = tuple(stash.pending_concurrency_keys)
@ -688,15 +704,18 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
if config is None:
return
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object")
if standard_logging_object is None:
standard_logging_object: Final = kwargs.get("standard_logging_object")
if not isinstance(standard_logging_object, dict):
return
# kwargs here is Logging.model_call_details, not the router's flat
# request kwargs admission sees: metadata/litellm_metadata are never
# top-level here, only nested under kwargs["litellm_params"] (see
# Logging.update_environment_variables).
litellm_params_for_metadata: Final = kwargs.get("litellm_params") or kwargs
litellm_params_raw: Final = kwargs.get("litellm_params")
litellm_params_for_metadata: Final[Mapping[str, object]] = (
litellm_params_raw if isinstance(litellm_params_raw, Mapping) else kwargs
)
metadata_variable_name: Final = _resolve_success_event_metadata_variable_name(litellm_params_for_metadata)
key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name)
key_alias: Final = _extract_key_alias(litellm_params_for_metadata, metadata_variable_name)
@ -761,7 +780,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
partition = await self._partition_for(partition_key) # not Final: rebound each loop iteration
accounting_task = asyncio.create_task( # not Final: rebound each loop iteration
partition.v3.async_increment_tokens_with_ttl_preservation(
pipeline_operations=tuple(group_operations), parent_otel_span=None
pipeline_operations=group_operations, parent_otel_span=None
)
)
_BACKGROUND_TASKS.add(accounting_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring

109
tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py
View file

@ -3,6 +3,9 @@ Unit tests for the global-scope, model-independent tag rate limiter.
"""
import asyncio
import os
import time
import uuid
from datetime import datetime, timedelta
import pytest
@ -48,6 +51,18 @@ def _data(tags: list[str], call_id: str = "call-1") -> dict:
return {"metadata": {"tags": tags}, "litellm_call_id": call_id}
def _redis_hook(time_controller: TimeController):
from litellm.caching.redis_cache import RedisCache
redis_host = os.getenv("REDIS_HOST")
redis_port = os.getenv("REDIS_PORT")
if not redis_host or not redis_port:
pytest.skip("Redis environment variables (REDIS_HOST, REDIS_PORT) not set")
redis_cache = RedisCache(host=redis_host, port=int(redis_port), password=os.getenv("REDIS_PASSWORD"))
dual_cache = DualCache(redis_cache=redis_cache)
return _PROXY_GlobalTagRateLimitsHook(internal_usage_cache=dual_cache, time_provider=time_controller.now), redis_cache
# ---------------------------------------------------------------------------
# No-op when unconfigured
# ---------------------------------------------------------------------------
@ -1400,3 +1415,97 @@ async def test_rejection_detail_does_not_disclose_the_resolved_tag_value(time_co
)
assert "tag_value" not in exc_info.value.detail
# ---------------------------------------------------------------------------
# Concurrency ttl refresh -- veria-ai finding: this hook's own atomic checks
# never carried refresh_ttl through to TAG_RL_CHECK_AND_INCR_SCRIPT or
# InMemoryCache.set_cache, even though it shares that script with
# model_based_tag_rate_limits_hook (whose own call sites got fixed first)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_redis_backed_concurrency_ttl_refreshes_on_every_admission(time_controller):
"""
A concurrency bucket isn't epoch-windowed like requests/tokens/dollars --
its ttl exists only as a crash-safety net for a reservation whose
explicit release never runs -- so a still-active bucket receiving
continuous admissions must keep extending that ttl, or it expires
mid-flight under sustained traffic, silently admitting past the cap.
"""
hook, redis_cache = _redis_hook(time_controller)
try:
await redis_cache.ping()
except Exception as e:
pytest.skip(f"Redis connection failed: {e!s}")
key = f"{{tag_rl:test:global-ttl-refresh:{uuid.uuid4().hex}}}:inflight"
cache = hook.internal_usage_cache
try:
admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True)
assert admitted
ttl_after_first_admission = await redis_cache.init_async_client().ttl(key)
assert ttl_after_first_admission > 0
await asyncio.sleep(2)
admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True)
assert admitted
ttl_after_second_admission = await redis_cache.init_async_client().ttl(key)
assert ttl_after_second_admission >= 2
finally:
await redis_cache.async_delete_cache(key=key)
@pytest.mark.asyncio
async def test_in_memory_concurrency_ttl_refreshes_on_every_admission(time_controller):
"""
In-memory mirror of the Redis test above: InMemoryCache.allow_ttl_override
leaves a still-live ttl untouched, so without refresh_ttl reaching
set_cache a concurrency counter's expiry stayed fixed from its first
admission even under sustained traffic.
"""
hook = _make_hook(time_controller)
cache = hook.internal_usage_cache
in_memory_cache = cache.dual_cache.in_memory_cache
key = f"tag_rl:test:global-in-memory-ttl-refresh:{uuid.uuid4().hex}"
admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True)
assert admitted
ttl_after_first_admission = in_memory_cache.ttl_dict[key]
admitted, _ = await hook._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True)
assert admitted
ttl_after_second_admission = in_memory_cache.ttl_dict[key]
assert ttl_after_second_admission > ttl_after_first_admission
@pytest.mark.asyncio
async def test_concurrency_limit_admission_refreshes_ttl_end_to_end(time_controller, monkeypatch):
"""
End-to-end regression through async_pre_call_hook itself (not just the
low-level _check_and_increment_one helper above): a concurrency entry's
bucket key must carry a live ttl after admission, proving refresh_ttl is
actually wired from the classified check through to the atomic batch,
not just present on the helper's own signature.
"""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 5, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
user_api_key_dict=_key(), cache=DualCache(), data=_data(["end_user_id:u1"]), call_type="completion"
)
in_memory_cache = hook.internal_usage_cache.dual_cache.in_memory_cache
inflight_keys = [key for key in in_memory_cache.ttl_dict if key.endswith(":inflight")]
assert len(inflight_keys) == 1
assert in_memory_cache.ttl_dict[inflight_keys[0]] > time.time()