tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 04:23:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #26954 from BerriAI/claude/lucid-margulis-e99b6b

Browse Source 
refactor(rate-limit): consolidate batch + dynamic limiter check/increment
This commit is contained in:
yuneng-jiang
2026-05-01 14:46:25 -07:00
committed by GitHub
parent ebd335da67 ca3e659a3c
commit dc681b9eb2
4 changed files with 1007 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/hooks/batch_rate_limiter.py
+27 -69
View File
@@ -164,13 +164,13 @@ class _PROXY_BatchRateLimiter(CustomLogger):
batch_usage: BatchFileUsage,
) -> None:
"""
Check rate limits and increment counters by the batch amounts.
Atomically check + increment rate-limit counters by the batch amounts.
Raises HTTPException if any limit would be exceeded.
Raises HTTPException if any descriptor would exceed its limit; in that
case no counter is modified. Backed by `atomic_check_and_increment_by_n`
which uses a Redis Lua script when available (multi-process atomic) and
falls back to a per-process asyncio.Lock + in-memory operation.
"""
from litellm.types.caching import RedisPipelineIncrementOperation
# Create descriptors and check if batch would exceed limits
descriptors = self.parallel_request_limiter._create_rate_limit_descriptors(
user_api_key_dict=user_api_key_dict,
data=data,
@@ -179,73 +179,31 @@ class _PROXY_BatchRateLimiter(CustomLogger):
model_has_failures=False,
)
# Check current usage without incrementing
rate_limit_response = await self.parallel_request_limiter.should_rate_limit(
descriptors=descriptors,
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=True,
)
increment: Dict[Literal["requests", "tokens"], int] = {
"requests": batch_usage.request_count,
"tokens": batch_usage.total_tokens,
}
increments: List[Dict[Literal["requests", "tokens"], int]] = [
increment for _ in descriptors
]
# Verify batch won't exceed any limits
for status in rate_limit_response["statuses"]:
rate_limit_type = status["rate_limit_type"]
limit_remaining = status["limit_remaining"]
required_capacity = (
batch_usage.request_count
if rate_limit_type == "requests"
else batch_usage.total_tokens if rate_limit_type == "tokens" else 0
)
if required_capacity > limit_remaining:
self._raise_rate_limit_error(
status, descriptors, batch_usage, rate_limit_type
)
# Build pipeline operations for batch increments
# Reuse the same keys that descriptors check
pipeline_operations: List[RedisPipelineIncrementOperation] = []
for descriptor in descriptors:
key = descriptor["key"]
value = descriptor["value"]
rate_limit = descriptor.get("rate_limit")
if rate_limit is None:
continue
# Add RPM increment if limit is set
if rate_limit.get("requests_per_unit") is not None:
rpm_key = self.parallel_request_limiter.create_rate_limit_keys(
key=key, value=value, rate_limit_type="requests"
)
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=rpm_key,
increment_value=batch_usage.request_count,
ttl=self.parallel_request_limiter.window_size,
)
)
# Add TPM increment if limit is set
if rate_limit.get("tokens_per_unit") is not None:
tpm_key = self.parallel_request_limiter.create_rate_limit_keys(
key=key, value=value, rate_limit_type="tokens"
)
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=tpm_key,
increment_value=batch_usage.total_tokens,
ttl=self.parallel_request_limiter.window_size,
)
)
# Execute increments
if pipeline_operations:
await self.parallel_request_limiter.async_increment_tokens_with_ttl_preservation(
pipeline_operations=pipeline_operations,
rate_limit_response = (
await self.parallel_request_limiter.atomic_check_and_increment_by_n(
descriptors=descriptors,
increments=increments,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
)
if rate_limit_response["overall_code"] == "OVER_LIMIT":
for status in rate_limit_response["statuses"]:
if status["code"] == "OVER_LIMIT":
self._raise_rate_limit_error(
status,
descriptors,
batch_usage,
status["rate_limit_type"],
)
async def count_input_file_usage(
self,
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
+109 -73
View File
@@ -4,7 +4,7 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting
import os
from datetime import datetime
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
from typing import TYPE_CHECKING, Callable, Dict, List, Literal, Optional, Union
from fastapi import HTTPException
@@ -460,92 +460,128 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
if priority_descriptors:
descriptors_to_check.extend(priority_descriptors)
# PHASE 1: Read-only check of ALL limits (no increments)
check_response = await self.v3_limiter.should_rate_limit(
descriptors=descriptors_to_check,
# Atomic check-and-increment for the ENFORCED descriptor set:
# - model_saturation_check is always enforced
# - priority_model is enforced only when saturation crosses threshold
#
# Backed by a Redis Lua script (multi-process atomic) with an
# asyncio.Lock + in-memory fallback for single-process deployments.
# All-or-nothing: if any enforced descriptor would exceed its limit,
# no counter is modified and the response carries "OVER_LIMIT".
enforced_descriptors: List[RateLimitDescriptor] = [model_wide_descriptor]
if priority_descriptors and should_enforce_priority:
enforced_descriptors.extend(priority_descriptors)
per_request_increment: Dict[Literal["requests", "tokens"], int] = {
"requests": 1,
"tokens": 0,
}
atomic_response = await self.v3_limiter.atomic_check_and_increment_by_n(
descriptors=enforced_descriptors,
increments=[per_request_increment for _ in enforced_descriptors],
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=True, # CRITICAL: Don't increment counters yet
)
verbose_proxy_logger.debug(
f"Read-only check: {json.dumps(check_response, indent=2)}"
f"Atomic check+increment response: {json.dumps(atomic_response, indent=2)}"
)
# PHASE 2: Decide which limits to enforce
if check_response["overall_code"] == "OVER_LIMIT":
for status in check_response["statuses"]:
if status["code"] == "OVER_LIMIT":
descriptor_key = status["descriptor_key"]
if atomic_response["overall_code"] == "OVER_LIMIT":
for status in atomic_response["statuses"]:
if status["code"] != "OVER_LIMIT":
continue
descriptor_key = status["descriptor_key"]
if descriptor_key == "model_saturation_check":
raise HTTPException(
status_code=429,
detail={
"error": f"Model capacity reached for {model}. "
f"Priority: {priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
"rate_limit_type": str(status["rate_limit_type"]),
"x-litellm-priority": priority or "default",
},
)
if descriptor_key == "priority_model":
verbose_proxy_logger.debug(
f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, "
f"priority: {priority}"
)
raise HTTPException(
status_code=429,
detail={
"error": f"Priority-based rate limit exceeded. "
f"Priority: {priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}, "
f"Model saturation: {saturation:.1%}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
"rate_limit_type": str(status["rate_limit_type"]),
"x-litellm-priority": priority or "default",
"x-litellm-saturation": f"{saturation:.2%}",
},
)
# Model-wide limit exceeded (ALWAYS enforce)
if descriptor_key == "model_saturation_check":
raise HTTPException(
status_code=429,
detail={
"error": f"Model capacity reached for {model}. "
f"Priority: {priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
"rate_limit_type": str(status["rate_limit_type"]),
"x-litellm-priority": priority or "default",
},
)
# Fail-closed guard: overall_code says OVER_LIMIT but no status
# matched a descriptor key we know how to translate into a 429.
# Refuse the request rather than silently fall through and let an
# over-limit request proceed to the model. Without this, a future
# caller wiring an unfamiliar descriptor into enforced_descriptors
# would silently bypass the rate limit.
offending = next(
(s for s in atomic_response["statuses"] if s["code"] == "OVER_LIMIT"),
None,
)
verbose_proxy_logger.error(
f"Dynamic rate limiter: OVER_LIMIT response with unknown "
f"descriptor_key(s) — refusing request. response={atomic_response}"
)
raise HTTPException(
status_code=429,
detail={
"error": "Rate limit exceeded",
"descriptor_key": (
offending["descriptor_key"] if offending else "unknown"
),
"rate_limit_type": (
str(offending["rate_limit_type"]) if offending else "unknown"
),
},
headers={
"retry-after": str(self.v3_limiter.window_size),
"x-litellm-priority": priority or "default",
},
)
# Priority limit exceeded (ONLY enforce when saturated)
elif descriptor_key == "priority_model" and should_enforce_priority:
verbose_proxy_logger.debug(
f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, "
f"priority: {priority}"
)
raise HTTPException(
status_code=429,
detail={
"error": f"Priority-based rate limit exceeded. "
f"Priority: {priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}, "
f"Model saturation: {saturation:.1%}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
"rate_limit_type": str(status["rate_limit_type"]),
"x-litellm-priority": priority or "default",
"x-litellm-saturation": f"{saturation:.2%}",
},
)
# PHASE 3: Increment counters separately to avoid early-exit issues
# Model counter must ALWAYS increment, but priority counter might be over limit
# If we increment them together, v3_limiter's in-memory check will exit early
# and skip incrementing the model counter
# Step 3a: Increment model-wide counter (always)
model_increment_response = await self.v3_limiter.should_rate_limit(
descriptors=[model_wide_descriptor],
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=False,
)
# Step 3b: Increment priority counter (may be over limit, but we still track it)
if priority_descriptors:
priority_increment_response = await self.v3_limiter.should_rate_limit(
# If priority is NOT enforced (saturation below threshold) but
# priority_descriptors exist, increment them for tracking only — no
# check, no rollback. This matches the prior tracking semantics.
#
# Using the non-atomic should_rate_limit (instead of
# atomic_check_and_increment_by_n) is intentional here: we don't want
# to enforce the limit, we only want to bump the counter so the
# priority allocation has accurate usage when it later becomes
# enforced. The increment-then-check semantics of should_rate_limit
# are fine because we ignore the OVER_LIMIT response.
if priority_descriptors and not should_enforce_priority:
priority_tracking_response = await self.v3_limiter.should_rate_limit(
descriptors=priority_descriptors,
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=False,
)
# Combine responses for post-call hook
combined_response = {
"overall_code": model_increment_response["overall_code"],
"statuses": model_increment_response["statuses"]
+ priority_increment_response["statuses"],
data["litellm_proxy_rate_limit_response"] = {
"overall_code": atomic_response["overall_code"],
"statuses": atomic_response["statuses"]
+ priority_tracking_response["statuses"],
}
data["litellm_proxy_rate_limit_response"] = combined_response
else:
data["litellm_proxy_rate_limit_response"] = model_increment_response
data["litellm_proxy_rate_limit_response"] = atomic_response
async def async_pre_call_hook(
self,
litellm/proxy/hooks/parallel_request_limiter_v3.py
+382
View File
@@ -4,6 +4,7 @@ This is a rate limiter implementation based on a similar one by Envoy proxy.
This is currently in development and not yet ready for production.
"""
import asyncio
import binascii
import os
from datetime import datetime
@@ -80,6 +81,90 @@ end
return results
"""
CHECK_AND_INCREMENT_BY_N_SCRIPT = """
-- Atomic check-and-increment-by-N across one or more descriptors.
-- All-or-nothing: if any descriptor would exceed its limit, no counter is
-- modified.
--
-- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied
-- timestamp so that window resets are deterministic across replicas with
-- skewed wall-clocks. This prevents a clock-skew-induced reopening of the
-- TOCTOU window across multi-replica deployments.
--
-- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor.
-- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]:
-- ARGV[(i-1)*4 + 1] = limit
-- ARGV[(i-1)*4 + 2] = increment
-- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets)
-- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length)
--
-- Return on success: { 0, new_counter_1, new_counter_2, ... }
-- Return on over-limit: { 1, descriptor_index, current_counter, limit }
local time_reply = redis.call('TIME')
local now = tonumber(time_reply[1])
local descriptor_count = #KEYS / 2
-- Pass 1: read state, validate. Abort without writing if any over limit.
local descriptor_state = {}
for i = 1, descriptor_count do
local window_key = KEYS[(i - 1) * 2 + 1]
local counter_key = KEYS[(i - 1) * 2 + 2]
local arg_base = (i - 1) * 4 + 1
local limit = tonumber(ARGV[arg_base])
local increment = tonumber(ARGV[arg_base + 1])
local window_size = tonumber(ARGV[arg_base + 3])
local window_start = redis.call('GET', window_key)
local window_expired = (not window_start) or
((now - tonumber(window_start)) >= window_size)
local current_counter
if window_expired then
current_counter = 0
else
current_counter = tonumber(redis.call('GET', counter_key) or 0)
end
if current_counter + increment > limit then
return { 1, i, current_counter, limit }
end
descriptor_state[i] = { window_expired, current_counter }
end
-- Pass 2: all checks passed. Apply increments.
local results = { 0 }
for i = 1, descriptor_count do
local window_key = KEYS[(i - 1) * 2 + 1]
local counter_key = KEYS[(i - 1) * 2 + 2]
local arg_base = (i - 1) * 4 + 1
local increment = tonumber(ARGV[arg_base + 1])
local ttl = tonumber(ARGV[arg_base + 2])
local window_size = tonumber(ARGV[arg_base + 3])
local window_expired = descriptor_state[i][1]
if window_expired then
redis.call('SET', window_key, tostring(now))
redis.call('SET', counter_key, increment)
redis.call('EXPIRE', window_key, window_size)
if ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(results, increment)
else
local new_counter = redis.call('INCRBY', counter_key, increment)
local current_ttl = redis.call('TTL', counter_key)
if current_ttl == -1 and ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(results, new_counter)
end
end
return results
"""
TOKEN_INCREMENT_SCRIPT = """
local results = {}
@@ -162,15 +247,37 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
TOKEN_INCREMENT_SCRIPT
)
)
self.check_and_increment_by_n_script = (
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
CHECK_AND_INCREMENT_BY_N_SCRIPT
)
)
else:
self.batch_rate_limiter_script = None
self.token_increment_script = None
self.check_and_increment_by_n_script = None
self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60))
# Batch rate limiter (lazy loaded)
self._batch_rate_limiter: Optional[Any] = None
# Serializes multi-phase check+increment sequences (batch + dynamic
# limiters) within this process to close the TOCTOU window between
# read-only check and counter increment. Multi-replica deployments
# additionally rely on Redis Lua atomicity for cross-process safety.
#
# Coarse granularity: this single lock serializes ALL atomic check+
# increment operations across batch and dynamic limiters on this
# instance. A slow batch input-file fetch (which happens upstream of
# the lock) does not block here, but Redis Lua latency does. If
# contention shows up under load (visible as p99 latency spikes
# correlated with batch traffic), shard to a per-descriptor-key lock
# via a `weakref.WeakValueDictionary[str, asyncio.Lock]`. Punted as a
# follow-up because Lua dominates wall-time and the lock is held for
# one round-trip.
self._check_and_increment_lock = asyncio.Lock()
def _get_batch_rate_limiter(self) -> Optional[Any]:
"""Get or lazy-load the batch rate limiter."""
if self._batch_rate_limiter is None:
@@ -588,6 +695,281 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
return rate_limit_response
async def atomic_check_and_increment_by_n(
self,
descriptors: List[RateLimitDescriptor],
increments: List[Dict[Literal["requests", "tokens"], int]],
parent_otel_span: Optional[Span] = None,
) -> RateLimitResponse:
"""
Atomic check-and-increment-by-N across one or more descriptors.
All-or-nothing: if any descriptor would exceed its limit, no counter is
modified and the response carries `overall_code = "OVER_LIMIT"` with
the offending descriptor's status. Closes the TOCTOU window between
read and increment in both single-process and multi-process (Redis)
deployments.
Args:
descriptors: rate-limit descriptors to check
increments: per-descriptor increment amounts, indexed parallel to
`descriptors`. Each entry is `{"requests": int, "tokens": int}`
— values default to 0 when a descriptor has no matching limit.
Returns:
RateLimitResponse with one status per (descriptor, rate_limit_type)
counter, mirroring `should_rate_limit`'s shape.
"""
if len(descriptors) != len(increments):
raise ValueError(
"atomic_check_and_increment_by_n: descriptors and increments "
"must have the same length"
)
keys: List[str] = []
per_counter_meta: List[Dict[str, Any]] = []
script_args: List[Any] = []
for descriptor, increment_amounts in zip(descriptors, increments):
descriptor_key = descriptor["key"]
descriptor_value = descriptor["value"]
rate_limit: RateLimitDescriptorRateLimitObject = (
descriptor.get("rate_limit") or RateLimitDescriptorRateLimitObject()
)
window_size = rate_limit.get("window_size") or self.window_size
window_key = f"{{{descriptor_key}:{descriptor_value}}}:window"
for rate_limit_type in ("requests", "tokens"):
rlt: Literal["requests", "tokens"] = cast(
Literal["requests", "tokens"], rate_limit_type
)
if rlt == "requests":
limit_value = rate_limit.get("requests_per_unit")
inc_amount = int(increment_amounts.get("requests", 0) or 0)
else:
limit_value = rate_limit.get("tokens_per_unit")
inc_amount = int(increment_amounts.get("tokens", 0) or 0)
if limit_value is None or inc_amount <= 0:
continue
counter_key = self.create_rate_limit_keys(
descriptor_key, descriptor_value, rlt
)
# Counter-key TTL and window_size are conceptually distinct
# ("how long the counter Redis key lives" vs "how long the
# sliding window is"). They happen to be equal today because
# we have no descriptor type that needs them apart, but they
# are kept as separate variables here so a future custom-TTL
# descriptor doesn't reintroduce a silent expiry bug. Both
# the Lua script and the in-memory fallback read these from
# their respective ARGV / meta slots.
ttl_seconds = int(window_size)
window_size_seconds = int(window_size)
keys.extend([window_key, counter_key])
# Per-counter 4-tuple matches the Lua ARGV layout exactly:
# [limit, increment, ttl_seconds, window_size_seconds].
script_args.extend(
[
int(limit_value),
inc_amount,
ttl_seconds,
window_size_seconds,
]
)
per_counter_meta.append(
{
"descriptor_key": descriptor_key,
"current_limit": int(limit_value),
"rate_limit_type": rlt,
"window_key": window_key,
"counter_key": counter_key,
"increment": inc_amount,
"ttl": ttl_seconds,
"window_size": window_size_seconds,
}
)
if not keys:
return RateLimitResponse(overall_code="OK", statuses=[])
# Multi-process atomicity via Redis Lua. Single-process atomicity
# falls back to the asyncio.Lock + in-memory sliding window below.
# Note: in-memory state diverges from Redis state — if Lua fails
# mid-write, retrying via in-memory may double-count. See fallback
# warning below.
if self.check_and_increment_by_n_script is not None:
try:
raw = await self.check_and_increment_by_n_script(
keys=keys,
args=script_args,
)
return self._build_atomic_response(raw, per_counter_meta)
except Exception as e:
# Escalated from warning to error: Lua failures (script timeout,
# Redis OOM, network partition) leave counter state ambiguous.
# The fallback path below uses LOCAL DualCache, which is a
# different store from Redis — counters here will diverge from
# Redis until that key's window expires (TTL bounds divergence).
# Operators should alert on this log line; sustained occurrences
# indicate Redis health degradation that may erode rate-limit
# accuracy.
verbose_proxy_logger.error(
f"atomic_check_and_increment_by_n: Redis Lua execution "
f"failed ({type(e).__name__}: {e}). Falling back to "
f"in-memory enforcement — counters will diverge from Redis "
f"state until window expires (window_size={self.window_size}s)."
)
async with self._check_and_increment_lock:
return await self._atomic_check_and_increment_in_memory(
per_counter_meta=per_counter_meta,
parent_otel_span=parent_otel_span,
)
def _build_atomic_response(
self,
raw: List[Any],
per_counter_meta: List[Dict[str, Any]],
) -> RateLimitResponse:
"""Convert Lua script return value to RateLimitResponse.
Indexing invariant: `per_counter_meta` and `KEYS` are parallel-indexed
at the COUNTER level, not the descriptor level. A descriptor with both
RPM and TPM limits emits two `(window_key, counter_key)` pairs and
two meta entries — one per counter. The Lua script's loop variable
`i` therefore enumerates counters, and the over-limit return tuple
`{1, i, ...}` carries a counter index that maps directly to
`per_counter_meta[i - 1]`. Keep these arrays parallel at the counter
level when modifying this code.
"""
if not raw:
return RateLimitResponse(overall_code="OK", statuses=[])
status_code = int(raw[0])
if status_code == 1:
# Over limit: { 1, counter_index (1-based), current_counter, limit }
descriptor_index = int(raw[1]) - 1
current_counter = int(raw[2])
limit = int(raw[3])
meta = per_counter_meta[descriptor_index]
return RateLimitResponse(
overall_code="OVER_LIMIT",
statuses=[
RateLimitStatus(
code="OVER_LIMIT",
current_limit=limit,
limit_remaining=max(0, limit - current_counter),
rate_limit_type=meta["rate_limit_type"],
descriptor_key=meta["descriptor_key"],
)
],
)
statuses: List[RateLimitStatus] = []
for meta, new_counter in zip(per_counter_meta, raw[1:]):
statuses.append(
RateLimitStatus(
code="OK",
current_limit=meta["current_limit"],
limit_remaining=max(0, meta["current_limit"] - int(new_counter)),
rate_limit_type=meta["rate_limit_type"],
descriptor_key=meta["descriptor_key"],
)
)
return RateLimitResponse(overall_code="OK", statuses=statuses)
async def _atomic_check_and_increment_in_memory(
self,
per_counter_meta: List[Dict[str, Any]],
parent_otel_span: Optional[Span] = None,
) -> RateLimitResponse:
"""In-memory all-or-nothing check-and-increment. Caller holds lock.
Reads/writes the LOCAL DualCache (`local_only=True`) — note this is
a different store from Redis. When this fallback fires after a Lua
failure, in-memory counters will diverge from Redis until each key's
window expires (TTL bounds divergence).
"""
# Use a single 'now' for the duration of this critical section so all
# descriptors evaluate window expiry consistently.
now_int = int(self._get_current_time().timestamp())
# Pass 1: read state, validate.
descriptor_state: List[Dict[str, Any]] = []
for meta in per_counter_meta:
window_size = meta["window_size"]
window_start = await self.internal_usage_cache.async_get_cache(
key=meta["window_key"],
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
window_expired = (
window_start is None or (now_int - int(window_start)) >= window_size
)
current_counter = (
0
if window_expired
else int(
await self.internal_usage_cache.async_get_cache(
key=meta["counter_key"],
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
or 0
)
)
if current_counter + meta["increment"] > meta["current_limit"]:
return RateLimitResponse(
overall_code="OVER_LIMIT",
statuses=[
RateLimitStatus(
code="OVER_LIMIT",
current_limit=meta["current_limit"],
limit_remaining=max(
0, meta["current_limit"] - current_counter
),
rate_limit_type=meta["rate_limit_type"],
descriptor_key=meta["descriptor_key"],
)
],
)
descriptor_state.append(
{"window_expired": window_expired, "current": current_counter}
)
# Pass 2: apply increments.
statuses: List[RateLimitStatus] = []
for meta, state in zip(per_counter_meta, descriptor_state):
new_counter = (
meta["increment"]
if state["window_expired"]
else state["current"] + meta["increment"]
)
if state["window_expired"]:
await self.internal_usage_cache.async_set_cache(
key=meta["window_key"],
value=str(now_int),
ttl=meta["window_size"],
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
await self.internal_usage_cache.async_set_cache(
key=meta["counter_key"],
value=new_counter,
ttl=meta["ttl"],
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
statuses.append(
RateLimitStatus(
code="OK",
current_limit=meta["current_limit"],
limit_remaining=max(0, meta["current_limit"] - new_counter),
rate_limit_type=meta["rate_limit_type"],
descriptor_key=meta["descriptor_key"],
)
)
return RateLimitResponse(overall_code="OK", statuses=statuses)
def create_organization_rate_limit_descriptor(
self, user_api_key_dict: UserAPIKeyAuth, requested_model: Optional[str] = None
) -> List[RateLimitDescriptor]:
tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py
+489
View File
@@ -0,0 +1,489 @@
"""
Tests validating TOCTOU race condition in batch + dynamic rate limiters.
Issue: rate-limit check (read_only=True) and counter increment happen as two
separate awaits. Concurrent requests all observe the same pre-increment state,
all pass validation, then all increment — bypassing the limit.
Vulnerable code paths:
- litellm/proxy/hooks/batch_rate_limiter.py:181-248
(_check_and_increment_batch_counters: should_rate_limit(read_only=True)
→ validate → async_increment_tokens_with_ttl_preservation)
- litellm/proxy/hooks/dynamic_rate_limiter_v3.py:463-548
(_check_rate_limits PHASE 1 read_only check → PHASE 3 increment)
These tests EXPECTED to fail against current (vulnerable) code and pass once
check-and-increment becomes atomic.
"""
import asyncio
import os
import sys
from typing import List
import pytest
sys.path.insert(0, os.path.abspath("../../../.."))
import litellm
from litellm import DualCache, Router
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.batch_rate_limiter import BatchFileUsage
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
_PROXY_DynamicRateLimitHandlerV3 as DynamicRateLimitHandler,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.utils import InternalUsageCache, hash_token
def _make_phase1_barrier(num_concurrent: int, timeout: float = 0.1):
"""
Sync primitive that, pre-fix, forces all N concurrent coroutines to finish
their read-only Phase 1 check before any proceeds to Phase 3 increment —
mimicking asyncio I/O scheduling under load on the vulnerable code.
Wraps `should_rate_limit` so on `read_only=True` calls it waits until N
callers arrive (TOCTOU window opened) OR `timeout` elapses (post-fix path:
the limiter's serialization lock prevents N from ever reaching the
barrier; the timeout lets the holder proceed so the lock can do its job).
Pre-fix: barrier fills before timeout → all see same state → bypass observed.
Post-fix: only lock-holder reaches barrier → times out → serial execution
enforces limit.
"""
arrived = 0
all_arrived = asyncio.Event()
def wrap(original):
async def patched(*args, **kwargs):
result = await original(*args, **kwargs)
if kwargs.get("read_only"):
nonlocal arrived
arrived += 1
if arrived >= num_concurrent:
all_arrived.set()
try:
await asyncio.wait_for(all_arrived.wait(), timeout=timeout)
except asyncio.TimeoutError:
pass
return result
return patched
return wrap
@pytest.mark.asyncio
async def test_batch_limiter_concurrent_bypasses_tpm_via_toctou():
"""
5 concurrent batch submissions of 40 tokens each against TPM=100 limit.
Sequential semantics: only 2 batches fit (2 * 40 = 80 ≤ 100, 3rd at 120 fails).
With TOCTOU: all 5 succeed → 200 tokens consumed, 100% over limit.
Demonstrates batch_rate_limiter.py:183-248 multi-phase flaw.
"""
NUM_CONCURRENT = 5
BATCH_TOKENS = 40
TPM_LIMIT = 100
dual_cache = DualCache()
internal_usage_cache = InternalUsageCache(dual_cache=dual_cache)
rate_limiter = _PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=internal_usage_cache
)
batch_limiter = rate_limiter._get_batch_rate_limiter()
assert batch_limiter is not None
user_api_key_dict = UserAPIKeyAuth(
api_key=hash_token("toctou-batch-key"),
tpm_limit=TPM_LIMIT,
rpm_limit=1000,
)
batch_usage = BatchFileUsage(total_tokens=BATCH_TOKENS, request_count=1)
barrier = _make_phase1_barrier(NUM_CONCURRENT)
rate_limiter.should_rate_limit = barrier(rate_limiter.should_rate_limit)
results = await asyncio.gather(
*[
batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=batch_usage,
)
for _ in range(NUM_CONCURRENT)
],
return_exceptions=True,
)
successes = [r for r in results if not isinstance(r, Exception)]
rejections = [r for r in results if isinstance(r, Exception)]
total_consumed = len(successes) * BATCH_TOKENS
max_allowed_successes = TPM_LIMIT // BATCH_TOKENS # 2
assert len(successes) <= max_allowed_successes, (
f"TOCTOU bypass: {len(successes)}/{NUM_CONCURRENT} concurrent batches "
f"passed despite TPM={TPM_LIMIT}. Consumed {total_consumed} tokens "
f"({total_consumed - TPM_LIMIT} over limit). "
f"Atomic check-and-increment would allow ≤{max_allowed_successes}. "
f"Rejections: {len(rejections)}"
)
@pytest.mark.asyncio
async def test_batch_limiter_uses_atomic_check_and_increment():
"""
Regression test: batch limiter routes through
`atomic_check_and_increment_by_n` rather than the legacy two-phase
pattern (read_only=True check + separate async_increment_tokens_with_ttl_preservation).
Ensures future refactors don't reintroduce the TOCTOU window.
"""
dual_cache = DualCache()
internal_usage_cache = InternalUsageCache(dual_cache=dual_cache)
rate_limiter = _PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=internal_usage_cache
)
batch_limiter = rate_limiter._get_batch_rate_limiter()
assert batch_limiter is not None
call_log: List[str] = []
original_atomic = rate_limiter.atomic_check_and_increment_by_n
original_should = rate_limiter.should_rate_limit
async def logging_atomic(*args, **kwargs):
call_log.append("atomic_check_and_increment_by_n")
return await original_atomic(*args, **kwargs)
async def logging_should(*args, **kwargs):
call_log.append(f"should_rate_limit(read_only={kwargs.get('read_only')})")
return await original_should(*args, **kwargs)
rate_limiter.atomic_check_and_increment_by_n = logging_atomic
rate_limiter.should_rate_limit = logging_should
user_api_key_dict = UserAPIKeyAuth(
api_key=hash_token("atomic-test-key"),
tpm_limit=10000,
rpm_limit=1000,
)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=BatchFileUsage(total_tokens=50, request_count=1),
)
assert "atomic_check_and_increment_by_n" in call_log, (
f"Batch limiter must route through atomic_check_and_increment_by_n. "
f"Calls observed: {call_log}"
)
legacy_calls = [c for c in call_log if c.startswith("should_rate_limit(")]
assert not legacy_calls, (
f"Batch limiter must not call should_rate_limit directly (legacy "
f"two-phase pattern). Observed: {legacy_calls}"
)
@pytest.mark.asyncio
async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity():
"""
DynamicRateLimitHandler PHASE 1 (read_only check) → PHASE 3 (increment)
is non-atomic: dynamic_rate_limiter_v3.py:463-548.
With TPM=100 model capacity and 5 concurrent priority="high" requests
each consuming the full model_saturation_check counter, all observe the
same Phase 1 state (counter=0), all pass, all proceed to Phase 3.
Sequential atomic semantics would block requests once the model counter
reaches its limit. TOCTOU lets all pass Phase 1 simultaneously.
"""
NUM_CONCURRENT = 10
MODEL_RPM = 2
# Sequential bound: dynamic limiter rejects when `counter > current_limit`
# (strict `>`), so a request whose Phase 1 sees counter=RPM still passes
# (RPM is not strictly greater). Atomic execution therefore admits up to
# RPM + 1 successes before the next sees counter > RPM.
MAX_SEQUENTIAL_SUCCESSES = MODEL_RPM + 1
os.environ["LITELLM_LICENSE"] = "test-license-key"
litellm.priority_reservation = {"high": 0.9, "low": 0.1}
dual_cache = DualCache()
handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)
model = "toctou-dyn-model"
llm_router = Router(
model_list=[
{
"model_name": model,
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": "test-key",
"api_base": "test-base",
"rpm": MODEL_RPM,
},
}
]
)
handler.update_variables(llm_router=llm_router)
barrier = _make_phase1_barrier(NUM_CONCURRENT)
handler.v3_limiter.should_rate_limit = barrier(handler.v3_limiter.should_rate_limit)
from litellm.types.router import ModelGroupInfo
model_group_info = ModelGroupInfo(
model_group=model,
providers=["openai"],
rpm=MODEL_RPM,
tpm=None,
)
async def one_request(idx: int):
user = UserAPIKeyAuth(api_key=hash_token(f"dyn-key-{idx}"))
user.metadata = {"priority": "high"}
try:
await handler._check_rate_limits(
model=model,
model_group_info=model_group_info,
user_api_key_dict=user,
priority="high",
saturation=0.0,
data={},
)
return "OK"
except Exception as e:
return e
results = await asyncio.gather(
*[one_request(i) for i in range(NUM_CONCURRENT)],
return_exceptions=True,
)
successes = [r for r in results if r == "OK"]
assert len(successes) <= MAX_SEQUENTIAL_SUCCESSES, (
f"TOCTOU bypass in DynamicRateLimitHandler: {len(successes)}/{NUM_CONCURRENT} "
f"concurrent requests passed Phase 1 + Phase 3 despite model RPM={MODEL_RPM}. "
f"Atomic check-and-increment would block once counter > RPM "
f"(at most {MAX_SEQUENTIAL_SUCCESSES} sequential successes)."
)
@pytest.mark.asyncio
async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment():
"""
Regression test: dynamic limiter's enforced descriptors flow through
`atomic_check_and_increment_by_n`, not the legacy
read_only=True check followed by a separate read_only=False increment.
When priority is enforced (saturation >= threshold), priority_model is
bundled into the atomic call alongside model_saturation_check. When not
enforced, priority counter is incremented for tracking only.
"""
os.environ["LITELLM_LICENSE"] = "test-license-key"
litellm.priority_reservation = {"high": 0.9, "low": 0.1}
dual_cache = DualCache()
handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)
model = "atomic-dyn-model"
llm_router = Router(
model_list=[
{
"model_name": model,
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": "test-key",
"api_base": "test-base",
"tpm": 1000,
},
}
]
)
handler.update_variables(llm_router=llm_router)
atomic_descriptors_observed: List[List[str]] = []
original_atomic = handler.v3_limiter.atomic_check_and_increment_by_n
async def logging_atomic(*args, **kwargs):
ds = kwargs.get("descriptors") or (args[0] if args else [])
atomic_descriptors_observed.append([d["key"] for d in ds])
return await original_atomic(*args, **kwargs)
handler.v3_limiter.atomic_check_and_increment_by_n = logging_atomic
from litellm.types.router import ModelGroupInfo
user = UserAPIKeyAuth(api_key=hash_token("dyn-atomic-key"))
user.metadata = {"priority": "high"}
await handler._check_rate_limits(
model=model,
model_group_info=ModelGroupInfo(
model_group=model,
providers=["openai"],
rpm=None,
tpm=1000,
),
user_api_key_dict=user,
priority="high",
saturation=0.0,
data={},
)
assert atomic_descriptors_observed, (
"Dynamic limiter must route enforced descriptors through "
"atomic_check_and_increment_by_n (no legacy read_only=True / "
"separate-increment pattern)."
)
assert "model_saturation_check" in atomic_descriptors_observed[0], (
f"Expected model_saturation_check in atomic descriptor set. "
f"Got: {atomic_descriptors_observed}"
)
@pytest.mark.asyncio
async def test_batch_zero_token_consumes_rpm_only():
"""
Zero-token batch (e.g. metadata-only call) should still increment RPM
counter but NOT TPM counter.
Edge case from review: `if inc_amount <= 0: continue` in
`atomic_check_and_increment_by_n` skips descriptor counters whose
increment is zero. Verifies asymmetric quota consumption is intentional
and observable: an RPM-bounded but TPM-free request path stays bounded
by RPM alone.
"""
dual_cache = DualCache()
internal_usage_cache = InternalUsageCache(dual_cache=dual_cache)
rate_limiter = _PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=internal_usage_cache
)
batch_limiter = rate_limiter._get_batch_rate_limiter()
assert batch_limiter is not None
user_api_key_dict = UserAPIKeyAuth(
api_key=hash_token("zero-token-key"),
tpm_limit=100,
rpm_limit=3,
)
zero_batch = BatchFileUsage(total_tokens=0, request_count=1)
# 3 zero-token batches must succeed (RPM=3 allows). 4th must hit RPM cap,
# NOT TPM (because token counter never increments past 0).
for i in range(3):
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=zero_batch,
)
# Inspect counters: RPM key incremented to 3, TPM key absent (or 0).
rpm_key = rate_limiter.create_rate_limit_keys(
"api_key", user_api_key_dict.api_key or "", "requests"
)
tpm_key = rate_limiter.create_rate_limit_keys(
"api_key", user_api_key_dict.api_key or "", "tokens"
)
rpm_val = await internal_usage_cache.async_get_cache(
key=rpm_key, litellm_parent_otel_span=None, local_only=True
)
tpm_val = await internal_usage_cache.async_get_cache(
key=tpm_key, litellm_parent_otel_span=None, local_only=True
)
assert int(rpm_val or 0) == 3, f"RPM counter must reach 3, got {rpm_val}"
assert tpm_val in (
None,
0,
"0",
), f"TPM counter must remain unset/0 for zero-token batches, got {tpm_val}"
# 4th attempt: RPM exhausted -> 429.
from fastapi import HTTPException
with pytest.raises(HTTPException) as exc:
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=zero_batch,
)
assert exc.value.status_code == 429
@pytest.mark.asyncio
async def test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor():
"""
Fail-closed guard: when atomic_check_and_increment_by_n returns
overall_code=OVER_LIMIT but with a descriptor_key the dispatcher does
not recognize, the dynamic limiter must raise 429 rather than silently
fall through.
Reproduces by patching atomic_check_and_increment_by_n to return an
OVER_LIMIT response carrying an unknown descriptor_key.
"""
from fastapi import HTTPException
os.environ["LITELLM_LICENSE"] = "test-license-key"
litellm.priority_reservation = {"high": 0.9, "low": 0.1}
dual_cache = DualCache()
handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)
model = "fail-closed-model"
llm_router = Router(
model_list=[
{
"model_name": model,
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": "test-key",
"api_base": "test-base",
"tpm": 1000,
},
}
]
)
handler.update_variables(llm_router=llm_router)
async def fake_atomic(*args, **kwargs):
return {
"overall_code": "OVER_LIMIT",
"statuses": [
{
"code": "OVER_LIMIT",
"current_limit": 100,
"limit_remaining": 0,
"rate_limit_type": "tokens",
"descriptor_key": "future_unrecognized_descriptor",
}
],
}
handler.v3_limiter.atomic_check_and_increment_by_n = fake_atomic
from litellm.types.router import ModelGroupInfo
user = UserAPIKeyAuth(api_key=hash_token("fail-closed-key"))
user.metadata = {"priority": "high"}
with pytest.raises(HTTPException) as exc:
await handler._check_rate_limits(
model=model,
model_group_info=ModelGroupInfo(
model_group=model,
providers=["openai"],
rpm=None,
tpm=1000,
),
user_api_key_dict=user,
priority="high",
saturation=0.0,
data={},
)
assert (
exc.value.status_code == 429
), f"Expected 429 fail-closed on unknown descriptor; got {exc.value.status_code}"