fix(proxy): enforce max_parallel_requests as a per-slot concurrency gauge (#32441)

* fix(proxy): enforce max_parallel_requests as a per-slot concurrency gauge

The v3 rate limiter tracked max_parallel_requests with the same
sliding-window machinery as RPM/TPM. A concurrency gauge cannot live on a
windowed counter: every window roll reset the counter to 1 while requests
were still in flight, the completion decrements for those forgotten
requests then drove the counter negative, and rejected requests left
stranded increments that nothing released. Under sustained load a key with
max_parallel_requests=5 let backend concurrency climb to the full client
concurrency (observed 60 on a live proxy) while the proxy kept returning
429s for everyone else

Replace the windowed counter with a per-slot registry (Redis sorted set of
slot ids scored by acquire time, with an asyncio-locked in-memory fallback):
admission atomically prunes expired slots and registers a new slot id only
when in_flight + 1 <= limit, so rejected requests never occupy a slot;
success, failure, and client-disconnect paths release exactly the slot id
this request acquired (stashed in the request metadata channels), so a
release without a matching acquire or a double-fired callback can never
free another request's slot; and a slot leaked by a crashed worker is
pruned individually after its TTL even under continuous traffic

Resolves LIT-4259
Fixes #16011

* fix(proxy): release every acquired gauge and respect mirrored counts in the in-memory fallback

Address review findings on the slot-registry gauge: the acquisition stash
now carries the gauge counter keys alongside the slot id, so the release
paths free the slot from every gauge it was registered under instead of
hardcoding the api_key scope, and the disconnect release keys off the
stashed acquisition instead of the key object's current
max_parallel_requests configuration (which can change mid-request). The
in-memory fallback now treats a cached integer (the count mirrored from
the last successful Redis script call) as real occupancy, carrying it
forward as a floored counter during a Redis outage instead of restarting
from an empty registry

* fix(proxy): release the parallel slot on proxy-level rejections

async_post_call_failure_hook is the only callback that fires when a
downstream hook (guardrail, budget check) rejects a request after the rate
limiter's pre-call hook acquired a slot; async_log_failure_event is a
completion-level callback and never runs for proxy-side rejections.
Release the stashed acquisition at the top of the hook, before the TPM
reservation guard, so those slots do not linger for the full slot TTL and
wedge the key at its limit under moderate rejection rates. Clearing the
acquisition marker keeps the release idempotent when a later failure
callback runs in the same flow

* test(proxy): cover success release, read-only count, Redis release mirror, and TPM rejection release

Four behaviors of the slot-registry gauge had no direct test: a successful
completion releasing exactly its acquired slot, read_only callers counting
in-flight slots through the count script (and degrading to the local
mirror when the script fails) without acquiring, the Redis release script
mirroring returned counts into the local cache, and the TPM reservation
rejection releasing the already-acquired slot before raising

* style(proxy): use builtin generics and union syntax in new rate limiter annotations

The slot-gauge code added Tuple/List/Dict and Optional[...] annotations, pushing
the UP006 and UP045 strict-rule totals past their ceilings in ruff-strict-budget.json.
Convert only the annotations this branch introduces to builtin generics and PEP 604
unions, leaving the rest of the module untouched.
This commit is contained in:
Yassin Kortam 2026-07-17 09:29:08 -07:00 • committed by GitHub
parent 637fc1f60e
commit 561b6796bc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 1242 additions and 173 deletions

View file

@ -2680,7 +2680,7 @@ class ProxyBaseLLMRequestProcessing:
# on disconnect, so the nested iterator hook (which only sees
# GeneratorExit on GC) cannot own the refund.
if not stream_completed:
proxy_logging_obj._release_max_parallel_requests_on_disconnect(user_api_key_dict)
proxy_logging_obj._release_max_parallel_requests_on_disconnect(user_api_key_dict, request_data)
client_disconnected = True
if not delivered_chunk:
from litellm.proxy.spend_tracking.budget_reservation import (

View file

@ -7,6 +7,7 @@ This is currently in development and not yet ready for production.
import asyncio
import binascii
import os
import uuid
from datetime import datetime
from typing import (
TYPE_CHECKING,
@ -185,6 +186,69 @@ end
return results
"""
PARALLEL_ACQUIRE_SCRIPT = """
-- Atomic check-and-acquire for the max_parallel_requests concurrency gauge.
-- Each gauge key is a sorted set of per-request slot ids scored by acquire
-- time (Redis server clock). In-flight requests are counted by ZCARD after
-- pruning slots older than the slot TTL, so unlike the windowed RPM/TPM
-- counters the gauge is never reset while requests are in flight, a
-- rejected request never occupies a slot, and a slot leaked by a crashed
-- worker self-heals after the slot TTL even under continuous traffic.
--
-- KEYS: one gauge zset key per descriptor.
-- ARGV: per-key triples (limit, slot_ttl_seconds, slot_id).
-- Success: { 0, in_flight_1, ... }. Over-limit: { 1, key_index, in_flight, limit }.
local time_reply = redis.call('TIME')
local now = tonumber(time_reply[1])
for i = 1, #KEYS do
local limit = tonumber(ARGV[(i - 1) * 3 + 1])
local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2])
redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - slot_ttl)
local in_flight = redis.call('ZCARD', KEYS[i])
if in_flight + 1 > limit then
return { 1, i, in_flight, limit }
end
end
local results = { 0 }
for i = 1, #KEYS do
local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2])
local slot_id = ARGV[(i - 1) * 3 + 3]
redis.call('ZADD', KEYS[i], now, slot_id)
redis.call('EXPIRE', KEYS[i], slot_ttl)
table.insert(results, redis.call('ZCARD', KEYS[i]))
end
return results
"""
PARALLEL_RELEASE_SCRIPT = """
-- Release one slot per gauge key by removing this request's slot id.
-- ZREM of an absent member (or key) is a no-op, so a release without a
-- matching acquire (proxy-side rejection, double-fired callback, slot
-- already expired) can never free a slot owned by another request.
-- KEYS: gauge zset keys. ARGV: per-key slot_id.
-- Returns the remaining in-flight count per key.
local results = {}
for i = 1, #KEYS do
redis.call('ZREM', KEYS[i], ARGV[i])
table.insert(results, redis.call('ZCARD', KEYS[i]))
end
return results
"""
PARALLEL_COUNT_SCRIPT = """
-- Read the current in-flight count per gauge key (prunes expired slots
-- first so leaked slots do not inflate the reading).
-- KEYS: gauge zset keys. ARGV: per-key slot_ttl_seconds.
local time_reply = redis.call('TIME')
local now = tonumber(time_reply[1])
local results = {}
for i = 1, #KEYS do
redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - tonumber(ARGV[i]))
table.insert(results, redis.call('ZCARD', KEYS[i]))
end
return results
"""
TOKEN_INCREMENT_SCRIPT = """
local results = {}
@ -248,6 +312,19 @@ RATE_LIMIT_DESCRIPTORS_KEY = "_litellm_rate_limit_descriptors"
# mirror ``x-ratelimit-*`` headers into the SLP. Streaming exits
# common_request_processing before ``async_post_call_success_hook`` runs.
RATE_LIMIT_RESPONSE_KEY = "_litellm_proxy_rate_limit_response"
# Holds the acquisition the pre-call hook made for this request: the slot id
# plus the gauge counter keys it was registered under. The success/failure
# callbacks release only this exact acquisition: those callbacks also fire
# for requests rejected at pre-call (which never acquired a slot), and an
# id-less release would free a slot still owned by another in-flight request
# — every rejection would then raise effective concurrency above the
# configured limit.
MAX_PARALLEL_SLOT_ACQUIRED_KEY = "_litellm_max_parallel_slot_acquired"
# How long an acquired slot counts toward the in-flight total before it is
# considered leaked (worker crashed without any release callback firing) and
# pruned. Also the longest request duration the gauge can track: a request
# running longer than this stops occupying its slot.
PARALLEL_REQUEST_SLOT_TTL_SECONDS = 3600
# Stash keys live ONLY in metadata channels — never at the top level of the
# request body. Top-level keys are forwarded as body params to upstream
# providers, which reject unknown fields with 400/429 errors.
@ -258,6 +335,7 @@ _LITELLM_STASH_KEYS: Tuple[str, ...] = (
TPM_RESERVATION_RELEASED_KEY,
RATE_LIMIT_DESCRIPTORS_KEY,
RATE_LIMIT_RESPONSE_KEY,
MAX_PARALLEL_SLOT_ACQUIRED_KEY,
)
@ -274,6 +352,17 @@ class RateLimitDescriptor(TypedDict):
rate_limit: Optional[RateLimitDescriptorRateLimitObject]
class ParallelRequestGauge(TypedDict):
counter_key: str
limit: int
descriptor_key: str
class ParallelSlotAcquisition(TypedDict):
slot_id: str
counter_keys: list[str]
class RateLimitStatus(TypedDict):
code: str
current_limit: int
@ -310,10 +399,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self.check_and_increment_by_n_script = (
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(CHECK_AND_INCREMENT_BY_N_SCRIPT)
)
self.parallel_acquire_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
PARALLEL_ACQUIRE_SCRIPT
)
self.parallel_release_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
PARALLEL_RELEASE_SCRIPT
)
self.parallel_count_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
PARALLEL_COUNT_SCRIPT
)
else:
self.batch_rate_limiter_script = None
self.token_increment_script = None
self.check_and_increment_by_n_script = None
self.parallel_acquire_script = None
self.parallel_release_script = None
self.parallel_count_script = None
self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60))
@ -559,7 +660,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
counter_key = keys_to_fetch[i + 1]
counter_value = cache_values[i + 1]
requests_limit = key_metadata[window_key]["requests_limit"]
max_parallel_requests_limit = key_metadata[window_key]["max_parallel_requests_limit"]
tokens_limit = key_metadata[window_key]["tokens_limit"]
# Determine which limit to use for current_limit and limit_remaining
@ -568,9 +668,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if counter_key.endswith(":requests"):
current_limit = requests_limit
rate_limit_type = "requests"
elif counter_key.endswith(":max_parallel_requests"):
current_limit = max_parallel_requests_limit
rate_limit_type = "max_parallel_requests"
elif counter_key.endswith(":tokens"):
current_limit = tokens_limit
rate_limit_type = "tokens"
@ -694,6 +791,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
parent_otel_span: Optional[Span] = None,
read_only: bool = False,
skip_tpm_check: bool = False,
parallel_slot_id: str | None = None,
) -> RateLimitResponse:
"""
Check if any of the rate limit descriptors should be rate limited.
@ -710,15 +808,122 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
``reserve_tpm_tokens`` reservation path should set this to
avoid the +1-per-key Lua / in-memory increment double-charging
the tokens counter.
``max_parallel_requests`` descriptors are enforced by the dedicated
concurrency-gauge path (``_check_parallel_request_gauges``), never by
the windowed counters. The gauge phase must stay AFTER the windowed
check so a windowed rejection never strands an acquired slot; the
reverse order would leak one gauge slot per RPM/TPM rejection.
``parallel_slot_id`` names the slot an admission registers; callers
that enforce (not read_only) should pass the id they will later
release with — when omitted, a generated slot id is used and the slot
can only be reclaimed by TTL expiry.
"""
current_time = self._get_current_time()
now = current_time.timestamp()
now_int = int(now) # Convert to integer for Redis Lua script
# Collect all keys and their metadata upfront
keys_to_fetch, key_metadata, gauges = self._collect_windowed_keys_and_gauges(
descriptors=descriptors,
skip_tpm_check=skip_tpm_check,
)
windowed_response = RateLimitResponse(overall_code="OK", statuses=[])
if keys_to_fetch:
## CHECK IN-MEMORY CACHE
cache_values = await self.internal_usage_cache.async_batch_get_cache(
keys=keys_to_fetch,
parent_otel_span=parent_otel_span,
local_only=True,
)
if cache_values is not None:
rate_limit_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata)
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.internal_usage_cache.async_batch_get_cache(
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
if cache_values is None:
cache_values = []
for _ in keys_to_fetch:
cache_values.append(str(now_int) if _.endswith(":window") else 0)
elif self.batch_rate_limiter_script is not None:
# NORMAL MODE: Increment counters in Redis
# Group keys by hash tag for Redis cluster compatibility
cache_values = await self._execute_redis_batch_rate_limiter_script(
keys_to_fetch=keys_to_fetch,
now_int=now_int,
)
# update in-memory cache with new values
for i in range(0, len(cache_values), 2):
window_key = keys_to_fetch[i]
counter_key = keys_to_fetch[i + 1]
window_value = cache_values[i]
counter_value = cache_values[i + 1]
await self.internal_usage_cache.async_set_cache(
key=counter_key,
value=counter_value,
ttl=self.window_size,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
await self.internal_usage_cache.async_set_cache(
key=window_key,
value=window_value,
ttl=self.window_size,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
else:
# NORMAL MODE: In-memory sliding window (no Redis)
cache_values = await self.in_memory_cache_sliding_window(
keys=keys_to_fetch,
now_int=now_int,
window_size=self.window_size,
)
windowed_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata)
if windowed_response["overall_code"] == "OVER_LIMIT":
return windowed_response
if not gauges:
return windowed_response
gauge_response = await self._check_parallel_request_gauges(
gauges=gauges,
slot_id=parallel_slot_id or uuid.uuid4().hex,
parent_otel_span=parent_otel_span,
read_only=read_only,
)
return RateLimitResponse(
overall_code=gauge_response["overall_code"],
statuses=[*windowed_response["statuses"], *gauge_response["statuses"]],
)
def _collect_windowed_keys_and_gauges(
self,
descriptors: list[RateLimitDescriptor],
skip_tpm_check: bool,
) -> tuple[list[str], dict[str, dict[str, Any]], list[ParallelRequestGauge]]:
"""
Split descriptors into the windowed (window_key, counter_key) fetch
list with its per-window metadata, and the concurrency gauges for
descriptors carrying a max_parallel_requests limit.
"""
keys_to_fetch: List[str] = []
key_metadata = {} # Store metadata for each key
key_metadata: dict[str, dict[str, Any]] = {}
gauges: list[ParallelRequestGauge] = []
for descriptor in descriptors:
descriptor_key = descriptor["key"]
descriptor_value = descriptor["value"]
@ -732,6 +937,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
window_key = f"{{{descriptor_key}:{descriptor_value}}}:window"
if max_parallel_requests_limit is not None:
gauges.append(
ParallelRequestGauge(
counter_key=self.create_rate_limit_keys(
descriptor_key, descriptor_value, "max_parallel_requests"
),
limit=int(max_parallel_requests_limit),
descriptor_key=descriptor_key,
)
)
rate_limit_set = False
if requests_limit is not None:
rpm_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, "requests")
@ -741,12 +957,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
tpm_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, "tokens")
keys_to_fetch.extend([window_key, tpm_key])
rate_limit_set = True
if max_parallel_requests_limit is not None:
max_parallel_requests_key = self.create_rate_limit_keys(
descriptor_key, descriptor_value, "max_parallel_requests"
)
keys_to_fetch.extend([window_key, max_parallel_requests_key])
rate_limit_set = True
if not rate_limit_set:
continue
@ -754,77 +964,252 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
key_metadata[window_key] = {
"requests_limit": (int(requests_limit) if requests_limit is not None else None),
"tokens_limit": int(tokens_limit) if tokens_limit is not None else None,
"max_parallel_requests_limit": (
int(max_parallel_requests_limit) if max_parallel_requests_limit is not None else None
),
"window_size": int(window_size),
"descriptor_key": descriptor_key,
}
return keys_to_fetch, key_metadata, gauges
## CHECK IN-MEMORY CACHE
cache_values = await self.internal_usage_cache.async_batch_get_cache(
keys=keys_to_fetch,
def _gauge_status(self, gauge: ParallelRequestGauge, in_flight: int, code: str) -> RateLimitStatus:
return RateLimitStatus(
code=code,
current_limit=gauge["limit"],
limit_remaining=max(0, gauge["limit"] - in_flight),
rate_limit_type="max_parallel_requests",
descriptor_key=gauge["descriptor_key"],
)
def _gauge_in_flight_from_cache_value(self, raw_value: Any) -> int:
"""
In-flight count from a cached gauge value: a dict of slot_id ->
acquire timestamp when the in-memory registry is authoritative, or
the mirrored integer count from the last Redis script result.
"""
if raw_value is None:
return 0
if isinstance(raw_value, dict):
cutoff = self._get_current_time().timestamp() - PARALLEL_REQUEST_SLOT_TTL_SECONDS
return sum(1 for ts in raw_value.values() if isinstance(ts, (int, float)) and ts >= cutoff)
return max(0, int(raw_value))
async def _check_parallel_request_gauges(
self,
gauges: list[ParallelRequestGauge],
slot_id: str,
parent_otel_span: Span | None = None,
read_only: bool = False,
) -> RateLimitResponse:
"""
Enforce max_parallel_requests as a concurrency gauge over a per-slot
registry: each admitted request registers ``slot_id`` with its
acquire time, and admission requires in_flight + 1 <= limit over the
unexpired slots. Unlike the windowed RPM/TPM counters, the gauge is
never reset while requests are in flight, a rejected request never
occupies a slot, and a slot leaked by a crashed worker is pruned
after PARALLEL_REQUEST_SLOT_TTL_SECONDS even under continuous
traffic. Releases remove exactly this request's slot id, so a
double-fired or unmatched release can never free another request's
slot.
"""
gauge_keys = [gauge["counter_key"] for gauge in gauges]
if read_only:
if self.parallel_count_script is not None:
try:
raw_counts = await self.parallel_count_script(
keys=gauge_keys,
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
verbose_proxy_logger.warning(f"parallel_count_script failed, using local mirror: {str(e)}")
counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span)
else:
counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span)
statuses = []
overall_code = "OK"
for gauge, in_flight in zip(gauges, counts):
code = "OVER_LIMIT" if in_flight >= gauge["limit"] else "OK"
if code == "OVER_LIMIT":
overall_code = "OVER_LIMIT"
statuses.append(self._gauge_status(gauge, in_flight, code))
return RateLimitResponse(overall_code=overall_code, statuses=statuses)
local_counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span)
for gauge, in_flight in zip(gauges, local_counts):
if in_flight >= gauge["limit"]:
return RateLimitResponse(
overall_code="OVER_LIMIT",
statuses=[self._gauge_status(gauge, in_flight, "OVER_LIMIT")],
)
if self.parallel_acquire_script is not None:
try:
raw = await self.parallel_acquire_script(
keys=gauge_keys,
args=[
arg for gauge in gauges for arg in (gauge["limit"], PARALLEL_REQUEST_SLOT_TTL_SECONDS, slot_id)
],
)
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to in-memory enforcement, never a 500
verbose_proxy_logger.warning(
f"parallel_acquire_script failed, falling back to in-memory gauge: {str(e)}"
)
async with self._check_and_increment_lock:
return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span)
if int(raw[0]) == 1:
gauge = gauges[int(raw[1]) - 1]
return RateLimitResponse(
overall_code="OVER_LIMIT",
statuses=[self._gauge_status(gauge, int(raw[2]), "OVER_LIMIT")],
)
statuses = []
for gauge, in_flight in zip(gauges, raw[1:]):
await self.internal_usage_cache.async_set_cache(
key=gauge["counter_key"],
value=int(in_flight),
ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
statuses.append(self._gauge_status(gauge, int(in_flight), "OK"))
return RateLimitResponse(overall_code="OK", statuses=statuses)
async with self._check_and_increment_lock:
return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span)
async def _read_local_gauge_counts(
self,
gauge_keys: list[str],
parent_otel_span: Span | None = None,
) -> list[int]:
values = await self.internal_usage_cache.async_batch_get_cache(
keys=gauge_keys,
parent_otel_span=parent_otel_span,
local_only=True,
)
if values is None:
return [0 for _ in gauge_keys]
return [self._gauge_in_flight_from_cache_value(value) for value in values]
if cache_values is not None:
rate_limit_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata)
if rate_limit_response["overall_code"] == "OVER_LIMIT":
return rate_limit_response
async def _acquire_parallel_slots_in_memory(
self,
gauges: list[ParallelRequestGauge],
slot_id: str,
parent_otel_span: Span | None = None,
) -> RateLimitResponse:
"""
All-or-nothing in-memory slot-registry acquire. Caller holds the lock.
## IF under limit in-memory, check Redis
if read_only:
# READ-ONLY MODE: Just read current values without incrementing
cache_values = await self.internal_usage_cache.async_batch_get_cache(
keys=keys_to_fetch,
parent_otel_span=parent_otel_span,
local_only=False, # Check Redis too
A cached dict is the authoritative in-memory registry. A cached
integer is the count mirrored from the last successful Redis script
call: when Redis fails over to this path, that mirror still counts
the slots in flight on the Redis side, so it is carried forward as
an integer counter (not discarded as an empty registry, which would
briefly double the admitted concurrency during a Redis outage).
"""
now = self._get_current_time().timestamp()
cutoff = now - PARALLEL_REQUEST_SLOT_TTL_SECONDS
states: list[tuple[dict[str, float] | None, int]] = []
for gauge in gauges:
raw_value = await self.internal_usage_cache.async_get_cache(
key=gauge["counter_key"],
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
if isinstance(raw_value, dict):
registry: dict[str, float] | None = {
key: float(ts) for key, ts in raw_value.items() if isinstance(ts, (int, float)) and ts >= cutoff
}
in_flight = len(registry or {})
elif raw_value is None:
registry = {}
in_flight = 0
else:
registry = None
in_flight = max(0, int(raw_value))
if in_flight + 1 > gauge["limit"]:
return RateLimitResponse(
overall_code="OVER_LIMIT",
statuses=[self._gauge_status(gauge, in_flight, "OVER_LIMIT")],
)
states.append((registry, in_flight))
# For keys that don't exist yet, set them to 0
if cache_values is None:
cache_values = []
for _ in keys_to_fetch:
cache_values.append(str(now_int) if _.endswith(":window") else 0)
elif self.batch_rate_limiter_script is not None:
# NORMAL MODE: Increment counters in Redis
# Group keys by hash tag for Redis cluster compatibility
cache_values = await self._execute_redis_batch_rate_limiter_script(
keys_to_fetch=keys_to_fetch,
now_int=now_int,
statuses = []
for gauge, (registry, in_flight) in zip(gauges, states):
new_value: Union[dict[str, float], int] = (
{**registry, slot_id: now} if registry is not None else in_flight + 1
)
await self.internal_usage_cache.async_set_cache(
key=gauge["counter_key"],
value=new_value,
ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
statuses.append(self._gauge_status(gauge, in_flight + 1, "OK"))
return RateLimitResponse(overall_code="OK", statuses=statuses)
# update in-memory cache with new values
for i in range(0, len(cache_values), 2):
window_key = keys_to_fetch[i]
counter_key = keys_to_fetch[i + 1]
window_value = cache_values[i]
counter_value = cache_values[i + 1]
async def _release_parallel_request_slots(
self,
acquisition: ParallelSlotAcquisition,
parent_otel_span: Span | None = None,
) -> None:
"""
Release the max_parallel_requests slots acquired at pre-call by
removing this request's slot id from every gauge it was registered
under. Removing an absent slot id is a no-op, so a release without a
matching acquire or a double-fired release can never free another
request's slot. The in-memory fallback decrements integer mirror
values (floored at 0) because the mirror carries no per-slot ids.
"""
counter_keys = acquisition["counter_keys"]
slot_id = acquisition["slot_id"]
if not counter_keys or not slot_id:
return
if self.parallel_release_script is not None:
try:
raw = await self.parallel_release_script(
keys=counter_keys,
args=[slot_id for _ in counter_keys],
)
for counter_key, remaining in zip(counter_keys, raw):
await self.internal_usage_cache.async_set_cache(
key=counter_key,
value=max(0, int(remaining)),
ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
return
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the in-memory release, never a 500
verbose_proxy_logger.warning(
f"parallel_release_script failed, falling back to in-memory release: {str(e)}"
)
async with self._check_and_increment_lock:
for counter_key in counter_keys:
raw_value = await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
if isinstance(raw_value, dict):
if slot_id not in raw_value:
continue
new_value: Union[dict[str, float], int] = {
key: ts for key, ts in raw_value.items() if key != slot_id
}
elif raw_value is None:
continue
else:
new_value = max(0, int(raw_value) - 1)
await self.internal_usage_cache.async_set_cache(
key=counter_key,
value=counter_value,
ttl=self.window_size,
value=new_value,
ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
await self.internal_usage_cache.async_set_cache(
key=window_key,
value=window_value,
ttl=self.window_size,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
else:
# NORMAL MODE: In-memory sliding window (no Redis)
cache_values = await self.in_memory_cache_sliding_window(
keys=keys_to_fetch,
now_int=now_int,
window_size=self.window_size,
)
rate_limit_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata)
return rate_limit_response
async def atomic_check_and_increment_by_n(
self,
@ -2027,10 +2412,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# shrinking the effective TPM budget by N and causing
# false-positive 429s under bursts. When reservation is disabled,
# this pass enforces TPM directly from the post-call counters.
parallel_counter_keys = [
self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests")
for d in descriptors
if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None
]
parallel_slot_id = uuid.uuid4().hex if parallel_counter_keys else None
response = await self.should_rate_limit(
descriptors=descriptors,
parent_otel_span=user_api_key_dict.parent_otel_span,
skip_tpm_check=self.tpm_reservation_enabled,
parallel_slot_id=parallel_slot_id,
)
if response["overall_code"] == "OVER_LIMIT":
@ -2049,6 +2442,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
key=RATE_LIMIT_RESPONSE_KEY,
value=response,
)
if parallel_slot_id is not None:
self._stash_value_in_metadata_channels(
data=data,
key=MAX_PARALLEL_SLOT_ACQUIRED_KEY,
value={
"slot_id": parallel_slot_id,
"counter_keys": parallel_counter_keys,
},
)
# ----------------------------------------------------------------
# TPM token reservation
@ -2108,6 +2510,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
if tpm_response["overall_code"] == "OVER_LIMIT":
acquisition = self._get_parallel_slot_acquisition(kwargs=data)
if acquisition is not None:
await self._release_parallel_request_slots(
acquisition=acquisition,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
self._clear_parallel_slot_marker(data)
self._handle_rate_limit_error(
response=tpm_response,
descriptors=descriptors,
@ -2480,6 +2889,50 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
"""True if a prior callback already refunded this request's reservation."""
return bool(cls._lookup_stashed_value(kwargs, standard_logging_metadata, TPM_RESERVATION_RELEASED_KEY))
@classmethod
def _get_parallel_slot_acquisition(
cls,
kwargs: Any,
standard_logging_metadata: dict[str, Any] | None = None,
) -> ParallelSlotAcquisition | None:
"""The slot acquisition this request's pre-call hook made, if any."""
candidate = cls._lookup_stashed_value(kwargs, standard_logging_metadata, MAX_PARALLEL_SLOT_ACQUIRED_KEY)
if not isinstance(candidate, dict):
return None
slot_id = candidate.get("slot_id")
counter_keys = candidate.get("counter_keys")
if not isinstance(slot_id, str) or not slot_id:
return None
if not isinstance(counter_keys, list) or not counter_keys:
return None
if not all(isinstance(key, str) and key for key in counter_keys):
return None
return ParallelSlotAcquisition(slot_id=slot_id, counter_keys=counter_keys)
@staticmethod
def _clear_parallel_slot_marker(data: Any) -> None:
"""
Remove the acquired-slot marker from every metadata channel a sibling
callback might read, so one release per acquire is an invariant even
when multiple callbacks fire for the same request.
"""
if not isinstance(data, dict):
return
for channel in ("metadata", "litellm_metadata"):
channel_dict = data.get(channel)
if isinstance(channel_dict, dict):
channel_dict.pop(MAX_PARALLEL_SLOT_ACQUIRED_KEY, None)
litellm_params = data.get("litellm_params")
if isinstance(litellm_params, dict):
lp_metadata = litellm_params.get("metadata")
if isinstance(lp_metadata, dict):
lp_metadata.pop(MAX_PARALLEL_SLOT_ACQUIRED_KEY, None)
slo = data.get("standard_logging_object")
if isinstance(slo, dict):
slo_meta = slo.get("metadata")
if isinstance(slo_meta, dict):
slo_meta.pop(MAX_PARALLEL_SLOT_ACQUIRED_KEY, None)
@staticmethod
def _mark_reservation_released(data: Any) -> None:
"""
@ -2621,7 +3074,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
standard_logging_object = kwargs.get("standard_logging_object") or {}
standard_logging_metadata = standard_logging_object.get("metadata") or {}
user_api_key = standard_logging_metadata.get("user_api_key_hash")
model_group = get_model_group_from_litellm_kwargs(kwargs)
# Get total tokens from response
@ -2658,20 +3110,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
pipeline_operations: List[RedisPipelineIncrementOperation] = []
# max_parallel_requests is its own counter (api-key only) — always decrement.
if user_api_key:
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=self.create_rate_limit_keys(
key="api_key",
value=user_api_key,
rate_limit_type="max_parallel_requests",
),
increment_value=-1,
ttl=self.window_size,
)
)
# ----------------------------------------------------------------
# TPM reconciliation
# Per-scope behavior:
@ -2719,6 +3157,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
try:
verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING")
standard_logging_object = kwargs.get("standard_logging_object") or {}
standard_logging_metadata = standard_logging_object.get("metadata") or {}
acquisition = self._get_parallel_slot_acquisition(
kwargs=kwargs,
standard_logging_metadata=standard_logging_metadata,
)
if acquisition is not None:
await self._release_parallel_request_slots(
acquisition=acquisition,
parent_otel_span=litellm_parent_otel_span,
)
self._clear_parallel_slot_marker(kwargs)
pipeline_operations = self._build_success_event_pipeline_operations(
kwargs=kwargs,
response_obj=response_obj,
@ -2855,22 +3306,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
litellm_parent_otel_span: Union[Span, None] = _get_parent_otel_span_from_kwargs(kwargs)
standard_logging_object = kwargs.get("standard_logging_object") or {}
standard_logging_metadata = standard_logging_object.get("metadata") or {}
user_api_key = standard_logging_metadata.get("user_api_key_hash")
pipeline_operations: List[RedisPipelineIncrementOperation] = []
if user_api_key:
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=self.create_rate_limit_keys(
key="api_key",
value=user_api_key,
rate_limit_type="max_parallel_requests",
),
increment_value=-1,
ttl=self.window_size,
)
acquisition = self._get_parallel_slot_acquisition(
kwargs=kwargs,
standard_logging_metadata=standard_logging_metadata,
)
if acquisition is not None:
await self._release_parallel_request_slots(
acquisition=acquisition,
parent_otel_span=litellm_parent_otel_span,
)
self._clear_parallel_slot_marker(kwargs)
# Skip the reservation refund if async_post_call_failure_hook
# already released it (proxy-level rejection that also bubbles up
@ -2920,40 +3368,35 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
except Exception as e:
verbose_proxy_logger.exception(f"Error in rate limit failure event: {str(e)}")
async def async_release_max_parallel_requests_on_disconnect(self, user_api_key_dict: UserAPIKeyAuth) -> None:
async def async_release_max_parallel_requests_on_disconnect(
self,
user_api_key_dict: UserAPIKeyAuth,
request_data: dict | None = None,
) -> None:
"""
Release the api-key ``max_parallel_requests`` slot that
``async_pre_call_hook`` reserved, for a request that ended without
``async_pre_call_hook`` acquired, for a request that ended without
either logging callback firing.
The +1 is normally undone by ``async_log_success_event`` (natural
The slot is normally released by ``async_log_success_event`` (natural
stream completion) or ``async_log_failure_event`` (LLM error). When a
client cancels a stream mid-flight, the cancellation surfaces as
``asyncio.CancelledError`` / ``GeneratorExit`` and neither callback
runs, so without this the counter leaks one slot per cancelled stream
until the key wedges at its limit.
runs, so without this the slot leaks per cancelled stream until its
TTL prunes it. ``request_data`` carries the stashed acquisition;
its presence (not the key object's current max_parallel_requests
configuration, which can change mid-request) decides whether there
is anything to release.
"""
if not user_api_key_dict.api_key or user_api_key_dict.max_parallel_requests is None:
acquisition = self._get_parallel_slot_acquisition(kwargs=request_data)
if acquisition is None:
return
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(
increment_list=[
RedisPipelineIncrementOperation(
key=self.create_rate_limit_keys(
key="api_key",
value=user_api_key_dict.api_key,
rate_limit_type="max_parallel_requests",
),
increment_value=-1,
# Refresh the window TTL on the decrement, matching the
# failure path. max_parallel_requests is a concurrency
# gauge, not a rolling-window count, so the key must
# outlive in-flight requests rather than expire mid-stream.
ttl=self.window_size,
)
],
litellm_parent_otel_span=None,
await self._release_parallel_request_slots(
acquisition=acquisition,
parent_otel_span=None,
)
self._clear_parallel_slot_marker(request_data)
async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response):
"""
@ -3002,17 +3445,29 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
traceback_str: Optional[str] = None,
) -> None:
"""
Release any TPM reservation when the request is rejected after the
pre-call hook reserved tokens but before the LLM call ran (e.g. a
downstream guardrail/auth hook raised). Without this, those
reservations are stranded — async_log_failure_event is a litellm
completion-level callback and never fires for proxy-side rejections.
Release the parallel-request slot and any TPM reservation when the
request is rejected after the pre-call hook acquired them but before
the LLM call ran (e.g. a downstream guardrail/auth hook raised).
Without this, those resources are stranded — async_log_failure_event
is a litellm completion-level callback and never fires for proxy-side
rejections, so a leaked slot would occupy the gauge for the full
PARALLEL_REQUEST_SLOT_TTL_SECONDS.
Idempotent via TPM_RESERVATION_RELEASED_KEY: if both this hook and
Idempotent: the slot release clears the acquisition marker (and slot
removal is a no-op ZREM on a second run), and the TPM refund is
guarded by TPM_RESERVATION_RELEASED_KEY — if both this hook and
async_log_failure_event end up running in the same flow, only the
first refund applies.
first release/refund applies.
"""
try:
acquisition = self._get_parallel_slot_acquisition(kwargs=request_data)
if acquisition is not None:
await self._release_parallel_request_slots(
acquisition=acquisition,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
self._clear_parallel_slot_marker(request_data)
if self._is_reservation_released(kwargs=request_data):
return
reserved_tokens = self._get_reserved_tokens_from_kwargs(kwargs=request_data)

View file

@ -7381,7 +7381,7 @@ async def async_data_generator(
# disconnect, so it fires reliably regardless of needs_iterator_wrap
# (a nested iterator hook would only see GeneratorExit on GC).
if not stream_completed:
proxy_logging_obj._release_max_parallel_requests_on_disconnect(user_api_key_dict)
proxy_logging_obj._release_max_parallel_requests_on_disconnect(user_api_key_dict, request_data)
client_disconnected = True
raise
except Exception as e:

View file

@ -2583,7 +2583,11 @@ class ProxyLogging:
logging_obj._deferred_stream_complete_args = None
asyncio.create_task(_deferred_cb(*_args))
def _release_max_parallel_requests_on_disconnect(self, user_api_key_dict: UserAPIKeyAuth) -> None:
def _release_max_parallel_requests_on_disconnect(
self,
user_api_key_dict: UserAPIKeyAuth,
request_data: dict | None = None,
) -> None:
"""
Release the api-key max_parallel_requests slot when a streaming
response is cancelled mid-flight (client disconnect). Neither the
@ -2603,14 +2607,16 @@ class ProxyLogging:
if not isinstance(limiter, _PROXY_MaxParallelRequestsHandler_v3):
return
try:
asyncio.create_task(limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict))
asyncio.create_task(
limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict, request_data)
)
except RuntimeError:
# No running event loop (e.g. interpreter/loop shutdown); the
# counter's window TTL will reclaim the slot.
verbose_proxy_logger.warning(
"parallel_request_limiter_v3: could not schedule "
"max_parallel_requests release on disconnect; no running "
"event loop. Slot will be reclaimed when its window TTL expires"
"event loop. Slot will be reclaimed when its TTL expires"
)
def _init_response_taking_too_long_task(self, data: Optional[dict] = None):

View file

@ -17,6 +17,10 @@ import litellm
from litellm import Router
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
MAX_PARALLEL_SLOT_ACQUIRED_KEY,
PARALLEL_REQUEST_SLOT_TTL_SECONDS,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler,
)
@ -566,10 +570,9 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_
# Verify that the correct token count was used based on the rate limit type
assert (
len(captured_operations) == 2
), "Should have 2 operations: max_parallel_requests decrement and TPM increment"
len(captured_operations) == 1
), "Should have 1 operation: the TPM increment (parallel slots are released via the gauge, not the pipeline)"
# Find the TPM increment operation (not the max_parallel_requests decrement)
tpm_operation = None
for op in captured_operations:
if op["key"].endswith(":tokens"):
@ -655,7 +658,10 @@ async def test_async_log_success_event_counts_non_chat_response_tokens(
@pytest.mark.asyncio
async def test_async_log_failure_event_v3():
"""
Simple test for async_log_failure_event - should decrement max_parallel_requests by 1
async_log_failure_event releases exactly this request's slot id: the
first release removes it, and repeated or unknown-slot releases are
no-ops that can never free another request's slot (releasing more than
was acquired is what previously let concurrency exceed the limit).
"""
_api_key = "sk-12345"
_api_key = hash_token(_api_key)
@ -663,33 +669,246 @@ async def test_async_log_failure_event_v3():
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
# Mock kwargs with user_api_key via standard_logging_object
mock_kwargs = {
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}
}
await _seed_max_parallel_requests_slots(local_cache, counter_key, ["slot-a", "slot-b"])
# Capture pipeline operations
captured_ops = []
def kwargs_with_slot(slot_id):
return {
"metadata": {
MAX_PARALLEL_SLOT_ACQUIRED_KEY: {
"slot_id": slot_id,
"counter_keys": [counter_key],
}
},
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
}
async def mock_pipeline(increment_list, **kwargs):
captured_ops.extend(increment_list)
async def in_flight():
return parallel_request_handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
)
parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = (
mock_pipeline
)
# Call async_log_failure_event
await parallel_request_handler.async_log_failure_event(
kwargs=mock_kwargs, response_obj=None, start_time=None, end_time=None
kwargs=kwargs_with_slot("slot-a"), response_obj=None, start_time=None, end_time=None
)
assert await in_flight() == 1
for slot_id in ("slot-a", "slot-unknown", "slot-a"):
await parallel_request_handler.async_log_failure_event(
kwargs=kwargs_with_slot(slot_id), response_obj=None, start_time=None, end_time=None
)
assert await in_flight() == 1
await parallel_request_handler.async_log_failure_event(
kwargs=kwargs_with_slot("slot-b"), response_obj=None, start_time=None, end_time=None
)
assert await in_flight() == 0
@pytest.mark.asyncio
async def test_failure_event_without_acquired_slot_does_not_release_v3():
"""
Failure callbacks also fire for requests rejected at pre-call, which never
acquired a parallel slot. Releasing on those frees a slot still owned by
another in-flight request, so every 429 would raise effective concurrency
above the configured limit. Without the acquired-slot marker the gauge
must stay untouched.
"""
_api_key = hash_token("sk-12345")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
await _seed_max_parallel_requests_slots(
local_cache, counter_key, ["slot-a", "slot-b", "slot-c"]
)
# Verify correct operation was created
assert len(captured_ops) == 1
op = captured_ops[0]
assert op["key"] == f"{{api_key:{_api_key}}}:max_parallel_requests"
assert op["increment_value"] == -1
assert op["ttl"] == 60 # default window size
await handler.async_log_failure_event(
kwargs={
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}
},
response_obj=None,
start_time=None,
end_time=None,
)
assert (
handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
)
== 3
)
@pytest.mark.asyncio
async def test_max_parallel_requests_not_reset_by_window_roll_v3():
"""
max_parallel_requests is a concurrency gauge, not a windowed counter: the
rate-limit window rolling over must not reset it while requests are still
in flight. Previously the gauge shared the sliding-window reset with
RPM/TPM, so every window roll forgot all in-flight requests and admitted
a fresh batch of `limit` on top of what was still running.
"""
controller = TimeController()
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache),
time_provider=controller.now,
)
_api_key = hash_token("sk-12345")
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=2)
for _ in range(2):
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-3.5-turbo"},
call_type="",
)
controller.advance(handler.window_size + 1)
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-3.5-turbo"},
call_type="",
)
assert exc_info.value.status_code == 429
assert "max_parallel_requests" in exc_info.value.detail
@pytest.mark.asyncio
async def test_rejected_request_does_not_consume_parallel_slot_v3():
"""
A 429-rejected request must not occupy a parallel-request slot: nothing
ever releases a slot for a request that was never admitted, so the old
increment-then-check behavior wedged the gauge above the limit and
rejected requests that should have been admitted after a release.
"""
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
_api_key = hash_token("sk-12345")
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1)
admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"}
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data=admitted_data,
call_type="",
)
acquisition = admitted_data["metadata"][MAX_PARALLEL_SLOT_ACQUIRED_KEY]
assert isinstance(acquisition, dict)
assert isinstance(acquisition["slot_id"], str) and acquisition["slot_id"]
assert acquisition["counter_keys"] == [f"{{api_key:{_api_key}}}:max_parallel_requests"]
for _ in range(3):
with pytest.raises(HTTPException):
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-3.5-turbo"},
call_type="",
)
await handler.async_log_failure_event(
kwargs={
"metadata": {MAX_PARALLEL_SLOT_ACQUIRED_KEY: acquisition},
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
},
response_obj=None,
start_time=None,
end_time=None,
)
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-3.5-turbo"},
call_type="",
)
@pytest.mark.asyncio
async def test_parallel_gauge_uses_atomic_redis_script_v3():
"""
With Redis available, gauge admission goes through the atomic
check-and-acquire script (limit, slot TTL, and this request's slot id as
args), the returned in-flight count is mirrored into the local cache,
and an over-limit script result maps to a 429 without occupying a slot.
"""
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
_api_key = hash_token("sk-12345")
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
captured_calls = []
async def fake_acquire(keys, args):
captured_calls.append((list(keys), list(args)))
return [0, 3]
handler.parallel_acquire_script = fake_acquire
data: Dict[str, Any] = {"model": "gpt-3.5-turbo"}
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data=data,
call_type="",
)
stashed_acquisition = data["metadata"][MAX_PARALLEL_SLOT_ACQUIRED_KEY]
assert isinstance(stashed_acquisition, dict)
stashed_slot_id = stashed_acquisition["slot_id"]
assert isinstance(stashed_slot_id, str) and stashed_slot_id
assert stashed_acquisition["counter_keys"] == [counter_key]
assert captured_calls == [
([counter_key], [5, PARALLEL_REQUEST_SLOT_TTL_SECONDS, stashed_slot_id])
]
assert (
await handler.internal_usage_cache.async_get_cache(
key=counter_key, litellm_parent_otel_span=None, local_only=True
)
== 3
)
gauge_statuses = [
s
for s in data["litellm_proxy_rate_limit_response"]["statuses"]
if s["rate_limit_type"] == "max_parallel_requests"
]
assert gauge_statuses == [
{
"code": "OK",
"current_limit": 5,
"limit_remaining": 2,
"rate_limit_type": "max_parallel_requests",
"descriptor_key": "api_key",
}
]
async def fake_acquire_over_limit(keys, args):
return [1, 1, 5, 5]
handler.parallel_acquire_script = fake_acquire_over_limit
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-3.5-turbo"},
call_type="",
)
assert exc_info.value.status_code == 429
assert "max_parallel_requests" in exc_info.value.detail
@pytest.mark.asyncio
@ -3227,27 +3446,28 @@ def test_get_key_mcp_rpm_limit_precedence():
assert get_team_mcp_rpm_limit(none_set) is None
async def _seed_max_parallel_requests_counter(
dual_cache: DualCache, counter_key: str, window_size: int
_TEST_SLOT_ID = "slot-disconnect-test"
async def _seed_max_parallel_requests_slots(
dual_cache: DualCache, counter_key: str, slot_ids: List[str]
) -> None:
await dual_cache.async_increment_cache_pipeline(
increment_list=[
RedisPipelineIncrementOperation(
key=counter_key, increment_value=1, ttl=window_size
)
]
await dual_cache.async_set_cache(
key=counter_key,
value={slot_id: time.time() for slot_id in slot_ids},
local_only=True,
)
async def _build_seeded_limiter():
"""Build a v3 limiter whose api-key counter already holds the pre-call +1."""
"""Build a v3 limiter whose api-key slot registry already holds the pre-call slot."""
api_key = hash_token("sk-disconnect")
cache = DualCache()
limiter = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache)
)
counter_key = f"{{api_key:{api_key}}}:max_parallel_requests"
await _seed_max_parallel_requests_counter(cache, counter_key, limiter.window_size)
await _seed_max_parallel_requests_slots(cache, counter_key, [_TEST_SLOT_ID])
user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2)
return limiter, cache, counter_key, user_api_key_dict
@ -3286,14 +3506,370 @@ async def test_release_max_parallel_requests_on_disconnect_v3():
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=2)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
await _seed_max_parallel_requests_counter(
local_cache, counter_key, handler.window_size
await _seed_max_parallel_requests_slots(local_cache, counter_key, [_TEST_SLOT_ID])
assert handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
) == 1
await handler.async_release_max_parallel_requests_on_disconnect(
user_api_key_dict,
request_data={
"metadata": {
MAX_PARALLEL_SLOT_ACQUIRED_KEY: {
"slot_id": _TEST_SLOT_ID,
"counter_keys": [counter_key],
}
}
},
)
assert await local_cache.async_get_cache(key=counter_key) == 1
await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
assert handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
) == 0
assert await local_cache.async_get_cache(key=counter_key) == 0
@pytest.mark.asyncio
async def test_release_on_disconnect_works_when_key_config_changed_v3():
"""
The disconnect release must be driven by the stashed acquisition, not the
key object's current max_parallel_requests configuration: if the limit is
cleared on the key while a request is in flight, the acquired slot still
has to be released or it lingers until TTL pruning.
"""
_api_key = hash_token("sk-12345")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
await _seed_max_parallel_requests_slots(local_cache, counter_key, [_TEST_SLOT_ID])
await handler.async_release_max_parallel_requests_on_disconnect(
UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=None),
request_data={
"metadata": {
MAX_PARALLEL_SLOT_ACQUIRED_KEY: {
"slot_id": _TEST_SLOT_ID,
"counter_keys": [counter_key],
}
}
},
)
assert handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
) == 0
@pytest.mark.asyncio
async def test_post_call_failure_hook_releases_parallel_slot_v3():
"""
A proxy-level rejection raised by a downstream hook after the rate
limiter's pre-call hook acquired a slot (guardrail, budget check) must
release that slot via async_post_call_failure_hook:
async_log_failure_event never fires for proxy-side rejections, so
without this the slot lingers for the full slot TTL and moderate
rejection rates wedge the key at its limit. The release must also be
idempotent with a later failure callback in the same flow.
"""
_api_key = hash_token("sk-12345")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"}
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data=admitted_data,
call_type="",
)
assert handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
) == 1
await handler.async_post_call_failure_hook(
request_data=admitted_data,
original_exception=Exception("guardrail rejected the request"),
user_api_key_dict=user_api_key_dict,
)
assert handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
) == 0
await handler.async_log_failure_event(
kwargs={
"metadata": admitted_data["metadata"],
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
},
response_obj=None,
start_time=None,
end_time=None,
)
assert handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
) == 0
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-3.5-turbo"},
call_type="",
)
@pytest.mark.asyncio
async def test_success_event_releases_parallel_slot_v3(monkeypatch):
"""
A successful completion must release exactly the slot its pre-call
acquired, freeing capacity for the next request; without it every
completed request would keep occupying the gauge until TTL pruning.
"""
_api_key = hash_token("sk-12345")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total")
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"}
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data=admitted_data,
call_type="",
)
assert handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
) == 1
await handler.async_log_success_event(
kwargs={
"metadata": admitted_data["metadata"],
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
},
response_obj=ModelResponse(
usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10)
),
start_time=datetime.now(),
end_time=datetime.now(),
)
assert handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
) == 0
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-3.5-turbo"},
call_type="",
)
@pytest.mark.asyncio
async def test_read_only_gauge_check_counts_without_acquiring_v3():
"""
read_only callers (e.g. the context-compaction pre-check) must observe
the in-flight count via the count script without registering a slot, and
a count-script failure must degrade to the local mirror instead of
raising.
"""
_api_key = hash_token("sk-12345")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
descriptors = [
{
"key": "api_key",
"value": _api_key,
"rate_limit": {"max_parallel_requests": 5},
}
]
captured_calls = []
async def fake_count(keys, args):
captured_calls.append((list(keys), list(args)))
return [3]
handler.parallel_count_script = fake_count
response = await handler.should_rate_limit(descriptors=descriptors, read_only=True)
assert captured_calls == [
([counter_key], [PARALLEL_REQUEST_SLOT_TTL_SECONDS])
]
assert response["overall_code"] == "OK"
assert response["statuses"] == [
{
"code": "OK",
"current_limit": 5,
"limit_remaining": 2,
"rate_limit_type": "max_parallel_requests",
"descriptor_key": "api_key",
}
]
assert await local_cache.async_get_cache(key=counter_key) is None
async def failing_count(keys, args):
raise ConnectionError("redis unavailable")
handler.parallel_count_script = failing_count
await _seed_max_parallel_requests_slots(
local_cache, counter_key, ["s1", "s2", "s3", "s4", "s5"]
)
response = await handler.should_rate_limit(descriptors=descriptors, read_only=True)
assert response["overall_code"] == "OVER_LIMIT"
assert response["statuses"][0]["rate_limit_type"] == "max_parallel_requests"
@pytest.mark.asyncio
async def test_redis_release_script_updates_local_mirror_v3():
"""
With Redis available, releases go through the release script with this
request's slot id per gauge key, and the returned in-flight counts are
mirrored into the local cache so the local first-pass check stays fresh.
"""
_api_key = hash_token("sk-12345")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
captured_calls = []
async def fake_release(keys, args):
captured_calls.append((list(keys), list(args)))
return [2]
handler.parallel_release_script = fake_release
await handler.async_log_failure_event(
kwargs={
"metadata": {
MAX_PARALLEL_SLOT_ACQUIRED_KEY: {
"slot_id": "slot-redis-test",
"counter_keys": [counter_key],
}
},
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
},
response_obj=None,
start_time=None,
end_time=None,
)
assert captured_calls == [([counter_key], ["slot-redis-test"])]
assert await local_cache.async_get_cache(key=counter_key) == 2
@pytest.mark.asyncio
async def test_tpm_over_limit_rejection_releases_parallel_slot_v3(monkeypatch):
"""
When the TPM reservation phase rejects a request AFTER the gauge slot was
acquired earlier in the same pre-call hook, the slot must be released
before the 429 is raised; otherwise every TPM rejection would leak a
slot until TTL pruning.
"""
monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
_api_key = hash_token("sk-12345")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
user_api_key_dict = UserAPIKeyAuth(
api_key=_api_key, max_parallel_requests=5, tpm_limit=100
)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
async def over_limit_reservation(descriptors, estimated_tokens, parent_otel_span=None):
return {
"overall_code": "OVER_LIMIT",
"statuses": [
{
"code": "OVER_LIMIT",
"current_limit": 100,
"limit_remaining": 0,
"rate_limit_type": "tokens",
"descriptor_key": "api_key",
}
],
}
monkeypatch.setattr(handler, "reserve_tpm_tokens", over_limit_reservation)
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hi"}]},
call_type="",
)
assert exc_info.value.status_code == 429
assert handler._gauge_in_flight_from_cache_value(
await local_cache.async_get_cache(key=counter_key)
) == 0
@pytest.mark.asyncio
async def test_in_memory_fallback_respects_mirrored_redis_count_v3():
"""
When Redis scripting fails after having worked, the local cache holds the
integer in-flight count mirrored from the last successful script call.
The in-memory fallback must treat that count as real occupancy (and
release must decrement it, floored at 0), not start over from an empty
registry, which would double the admitted concurrency during a Redis
outage.
"""
_api_key = hash_token("sk-12345")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
async def failing_script(keys, args):
raise ConnectionError("redis unavailable")
handler.parallel_acquire_script = failing_script
handler.parallel_release_script = failing_script
await local_cache.async_set_cache(key=counter_key, value=5, local_only=True)
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "gpt-3.5-turbo"},
call_type="",
)
assert exc_info.value.status_code == 429
await local_cache.async_set_cache(key=counter_key, value=4, local_only=True)
admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"}
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data=admitted_data,
call_type="",
)
assert await local_cache.async_get_cache(key=counter_key) == 5
await handler.async_log_failure_event(
kwargs={
"metadata": admitted_data["metadata"],
"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
},
response_obj=None,
start_time=None,
end_time=None,
)
assert await local_cache.async_get_cache(key=counter_key) == 4
@pytest.mark.asyncio
@ -3338,7 +3914,9 @@ async def test_async_streaming_data_generator_releases_counter_on_disconnect_v3(
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
limiter, cache, counter_key, user_api_key_dict = await _build_seeded_limiter()
assert await cache.async_get_cache(key=counter_key) == 1
assert limiter._gauge_in_flight_from_cache_value(
await cache.async_get_cache(key=counter_key)
) == 1
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = limiter
@ -3354,7 +3932,15 @@ async def test_async_streaming_data_generator_releases_counter_on_disconnect_v3(
gen = ProxyBaseLLMRequestProcessing.async_sse_data_generator(
response=upstream(),
user_api_key_dict=user_api_key_dict,
request_data={"model": "claude-test"},
request_data={
"model": "claude-test",
"metadata": {
MAX_PARALLEL_SLOT_ACQUIRED_KEY: {
"slot_id": _TEST_SLOT_ID,
"counter_keys": [counter_key],
}
},
},
proxy_logging_obj=proxy_logging_obj,
)
await gen.__anext__()
@ -3365,7 +3951,9 @@ async def test_async_streaming_data_generator_releases_counter_on_disconnect_v3(
await gen.aclose()
await _drain_release_task()
assert await cache.async_get_cache(key=counter_key) == 0
assert limiter._gauge_in_flight_from_cache_value(
await cache.async_get_cache(key=counter_key)
) == 0
@pytest.mark.parametrize("disconnect", ["cancel", "aclose"])
@ -3399,7 +3987,15 @@ async def test_async_data_generator_releases_counter_on_disconnect_v3(disconnect
gen = proxy_server.async_data_generator(
response=upstream(),
user_api_key_dict=user_api_key_dict,
request_data={"model": "gpt-test"},
request_data={
"model": "gpt-test",
"metadata": {
MAX_PARALLEL_SLOT_ACQUIRED_KEY: {
"slot_id": _TEST_SLOT_ID,
"counter_keys": [counter_key],
}
},
},
)
await gen.__anext__()
if disconnect == "cancel":
@ -3408,7 +4004,9 @@ async def test_async_data_generator_releases_counter_on_disconnect_v3(disconnect
else:
await gen.aclose()
await _drain_release_task()
assert await cache.async_get_cache(key=counter_key) == 0
assert limiter._gauge_in_flight_from_cache_value(
await cache.async_get_cache(key=counter_key)
) == 0
finally:
if saved_hook is not None:
proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = (
@ -3452,12 +4050,22 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3():
gen = proxy_server.async_data_generator(
response=upstream(),
user_api_key_dict=user_api_key_dict,
request_data={"model": "gpt-test"},
request_data={
"model": "gpt-test",
"metadata": {
MAX_PARALLEL_SLOT_ACQUIRED_KEY: {
"slot_id": _TEST_SLOT_ID,
"counter_keys": [counter_key],
}
},
},
)
await gen.__anext__()
await gen.aclose()
await _drain_release_task()
assert await cache.async_get_cache(key=counter_key) == 0
assert limiter._gauge_in_flight_from_cache_value(
await cache.async_get_cache(key=counter_key)
) == 0
finally:
if saved_hook is not None:
proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = (