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

Merge pull request #37539 from BerriAI/litellm_batch_enqueued_token_limit

Browse source 
feat(proxy): enqueued-token rate limiting for batches with refund on completion and cancellation
This commit is contained in:
Mateo Wang 2026-08-19 22:22:05 -07:00 • committed by GitHub
commit 47a7e1742e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 1862 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/constants.py
View file

@ -1766,6 +1766,17 @@ PTU_LAPSED_ALERT_LIMIT: Final[int] = 10
# one is seconds old, so a few minutes separates them.
PTU_PRUNE_SKEW_GRACE_SECONDS: Final[int] = 300
# How long enqueued-token reservations for batches live without a refund. Providers
# complete or expire batches within their completion window (24h for OpenAI), so a
# reservation still unrefunded after 8 days belongs to a batch whose terminal state
# was never observed (e.g. proxy restart); expiry returns the tokens to the caller.
BATCH_ENQUEUED_TOKEN_TTL_SECONDS: Final[int] = 8 * 24 * 60 * 60
# Key/team metadata field that opts batches into enqueued-token limiting. Only proxy
# admins may write it: when present it replaces the standard RPM/TPM checks for
# batch submissions.
BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit"
# Shared read-only empty mapping, for defaulting optional Mapping parameters without
# constructing a fresh mutable dict at each call site.
EMPTY_MAPPING: Final = MappingProxyType({})

47
litellm/proxy/auth/auth_utils.py
View file

@ -12,7 +12,12 @@ from pydantic import PositiveInt, TypeAdapter, ValidationError
import litellm
from litellm import Router, provider_list
from litellm._logging import verbose_proxy_logger
from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH, STANDARD_CUSTOMER_ID_HEADERS
from litellm.constants import (
BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY,
EMPTY_MAPPING,
MINIMUM_CUSTOM_KEY_LENGTH,
STANDARD_CUSTOMER_ID_HEADERS,
)
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.url_utils import (
SSRFError,
@ -1169,6 +1174,46 @@ def enforce_output_token_estimates_are_admin_only(
)
class BatchEnqueuedTokenLimitRequest(Protocol):
"""The shape of any management request that can carry a batch enqueued-token limit."""
@property
def metadata(self) -> Mapping[str, object] | None: ...
@property
def model_fields_set(self) -> Collection[str]: ...
def enforce_batch_enqueued_token_limit_is_admin_only(
data: BatchEnqueuedTokenLimitRequest,
existing_metadata: Mapping[str, object] | None,
user_api_key_dict: UserAPIKeyAuth,
entity: Literal["key", "team"],
) -> None:
"""Only a proxy admin may change a key or team's batch enqueued-token limit.
When set, ``batch_enqueued_token_limit`` replaces the standard RPM/TPM checks
for batch submissions, so a holder-writable copy would let a caller lift their
own batch quota. Gated on the resulting value rather than on presence, so a
form resending the stored value stays a no-op.
"""
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
stored: Final[Mapping[str, object]] = existing_metadata or EMPTY_MAPPING
requested: Final[Mapping[str, object]] = (
(data.metadata or EMPTY_MAPPING) if "metadata" in data.model_fields_set else stored
)
if requested.get(BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY) == stored.get(BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY):
return
raise HTTPException(
status_code=403,
detail={ # mutable-ok: HTTPException.detail has no immutable form
"error": f"Only proxy admins can set {BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY} on a {entity}. "
"It replaces the standard rate limit checks for batch submissions."
},
)
def get_model_rate_limit_from_metadata(
user_api_key_dict: UserAPIKeyAuth,
metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"],

456
litellm/proxy/hooks/batch_enqueued_tokens.py Normal file
View file

@ -0,0 +1,456 @@
"""
Enqueued-token accounting for batch submissions.
Opt-in via admin-set ``batch_enqueued_token_limit`` in key or team metadata: batch
submissions reserve their estimated token count against a long-lived
enqueued-token allowance instead of the per-minute rate-limit windows, and
the reservation is refunded when the batch reaches a terminal state
(completed, failed, expired, or cancelled).
"""
import asyncio
import math
import time
import uuid
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.constants import BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, BATCH_ENQUEUED_TOKEN_TTL_SECONDS
from litellm.proxy._types import UserAPIKeyAuth
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
Span = _Span
InternalUsageCache = _InternalUsageCache
BATCH_ENQUEUED_REFUND_STATUSES: Final[frozenset[str]] = frozenset(
{"completed", "complete", "failed", "expired", "cancelled", "cancelling"}
)
ScopeKey: TypeAlias = Literal["api_key", "team"]
RESERVE_ENQUEUED_TOKENS_SCRIPT: Final = """
local amount = tonumber(ARGV[1])
local ttl = tonumber(ARGV[2])
local limit = tonumber(ARGV[3])
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
if current + amount > limit then
return {0, current}
end
local updated = redis.call('INCRBY', KEYS[1], amount)
redis.call('EXPIRE', KEYS[1], ttl)
return {1, updated}
"""
REFUND_ENQUEUED_TOKENS_SCRIPT: Final = """
local updated = redis.call('DECRBY', KEYS[1], tonumber(ARGV[1]))
if updated <= 0 then
redis.call('DEL', KEYS[1])
end
return 1
"""
SAVE_RESERVATION_SCRIPT: Final = """
redis.call('SET', KEYS[1], ARGV[1], 'EX', tonumber(ARGV[2]))
return 1
"""
POP_RESERVATION_SCRIPT: Final = """
local value = redis.call('GET', KEYS[1])
if value and value ~= '' then
redis.call('SET', KEYS[1], '', 'EX', tonumber(ARGV[1]))
end
return value
"""
@dataclass(frozen=True, slots=True)
class BatchEnqueuedTokenScope:
key: ScopeKey
value: str
limit: int
ReservationBackend: TypeAlias = Literal["redis", "memory"]
@dataclass(frozen=True, slots=True)
class BatchEnqueuedTokenReservation:
tokens: int
scopes: tuple[BatchEnqueuedTokenScope, ...]
backend: ReservationBackend = "redis"
owner: str = ""
reserved_at_monotonic: float = field(default_factory=time.monotonic, compare=False)
@dataclass(frozen=True, slots=True)
class BatchEnqueuedTokenOverLimit:
scope: BatchEnqueuedTokenScope
enqueued: int
BatchEnqueuedTokenOutcome: TypeAlias = BatchEnqueuedTokenReservation | BatchEnqueuedTokenOverLimit
_LIMIT_ADAPTER: Final = TypeAdapter(Annotated[int, Field(gt=0)])
_RESERVE_RESULT_ADAPTER: Final = TypeAdapter(tuple[int, int])
_POPPED_VALUE_ADAPTER: Final = TypeAdapter(str | bytes | None)
_STORED_COUNTER_ADAPTER: Final = TypeAdapter(int | None)
_RESERVATION_ADAPTER: Final = TypeAdapter(BatchEnqueuedTokenReservation)
class _ScriptRunner(Protocol):
def __call__(self, keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> Awaitable[object]: ...
def _read_metadata_limit(metadata: Mapping[str, object] | None) -> int | None:
if not metadata:
return None
raw: Final = metadata.get(BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY)
if raw is None:
return None
try:
return _LIMIT_ADAPTER.validate_python(raw)
except ValidationError:
verbose_proxy_logger.warning(
"Ignoring invalid %s value %r; expected a positive integer",
BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY,
raw,
)
return None
def resolve_batch_enqueued_token_scopes(
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[BatchEnqueuedTokenScope, ...]:
key_limit: Final = _read_metadata_limit(user_api_key_dict.metadata)
team_limit: Final = _read_metadata_limit(user_api_key_dict.team_metadata)
candidates: Final = (
BatchEnqueuedTokenScope(key="api_key", value=user_api_key_dict.api_key, limit=key_limit)
if key_limit is not None and user_api_key_dict.api_key
else None,
BatchEnqueuedTokenScope(key="team", value=user_api_key_dict.team_id, limit=team_limit)
if team_limit is not None and user_api_key_dict.team_id
else None,
)
return tuple(scope for scope in candidates if scope is not None)
def canonical_provider_batch_id(batch_id: str) -> str:
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id, # pyright: ignore[reportPrivateUsage] # canonical unified-id decoder has no public wrapper
get_batch_id_from_unified_batch_id,
get_original_file_id,
)
decoded: Final = _is_base64_encoded_unified_file_id(batch_id)
if isinstance(decoded, str):
if "llm_batch_id" in decoded or "generic_response_id" in decoded:
return get_batch_id_from_unified_batch_id(decoded)
return decoded
return get_original_file_id(batch_id)
class _BatchResponseView(BaseModel):
model_config = ConfigDict(extra="ignore")
id: str
status: str
object: Literal["batch"]
def batch_response_view(response: object) -> _BatchResponseView | None:
try:
return _BatchResponseView.model_validate(response, from_attributes=True)
except ValidationError:
return None
class BatchEnqueuedTokenStore:
"""Tracks enqueued batch tokens per scope, plus per-batch reservation records for refunds.
Counters and records live in Redis when Redis is configured, through
single-key Lua scripts issued one scope at a time (Redis Cluster safe: no
cross-slot commands), with an over-limit or failing scope rolling back the
scopes reserved before it; otherwise a single-process in-memory fallback
guarded by one asyncio lock is used. Reservations remember which backend
granted them, and in-memory grants also remember the granting worker, so a
refund never debits counters the grant did not charge. Everything expires after
``BATCH_ENQUEUED_TOKEN_TTL_SECONDS`` so a crash between submission and the
terminal-state refund can never leak tokens forever, and reservation records
expire no later than the counters they would refund, so a stale record can
never debit an allowance re-granted after its counters expired.
"""
def __init__(
self,
internal_usage_cache: "InternalUsageCache",
monotonic: Callable[[], float] = time.monotonic,
) -> None:
self.internal_usage_cache = internal_usage_cache
self._monotonic: Final = monotonic
self._lock = asyncio.Lock()
self._owner_token = uuid.uuid4().hex
redis_cache = internal_usage_cache.dual_cache.redis_cache
self._reserve_script: _ScriptRunner | None = (
redis_cache.async_register_script(RESERVE_ENQUEUED_TOKENS_SCRIPT) if redis_cache is not None else None
)
self._refund_script: _ScriptRunner | None = (
redis_cache.async_register_script(REFUND_ENQUEUED_TOKENS_SCRIPT) if redis_cache is not None else None
)
self._save_script: _ScriptRunner | None = (
redis_cache.async_register_script(SAVE_RESERVATION_SCRIPT) if redis_cache is not None else None
)
self._pop_script: _ScriptRunner | None = (
redis_cache.async_register_script(POP_RESERVATION_SCRIPT) if redis_cache is not None else None
)
@staticmethod
def _counter_key(scope: BatchEnqueuedTokenScope) -> str:
return f"batch_enqueued_tokens:{scope.key}:{scope.value}"
@staticmethod
def _record_key(batch_id: str) -> str:
return f"batch_enqueued_token_reservation:{batch_id}"
async def reserve(
self,
tokens: int,
scopes: tuple[BatchEnqueuedTokenScope, ...],
litellm_parent_otel_span: "Span | None" = None,
) -> BatchEnqueuedTokenOutcome:
if tokens <= 0 or not scopes:
return BatchEnqueuedTokenReservation(tokens=max(tokens, 0), scopes=scopes)
reserve_script: Final = self._reserve_script
refund_script: Final = self._refund_script
if reserve_script is not None and refund_script is not None:
try:
return await self._reserve_via_redis(reserve_script, refund_script, tokens=tokens, scopes=scopes)
except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory counters
verbose_proxy_logger.warning(
"Redis enqueued-token reserve failed, falling back to in-memory: %s", str(e)
)
return await self._reserve_in_memory(tokens=tokens, scopes=scopes, span=litellm_parent_otel_span)
async def _reserve_via_redis(
self,
reserve_script: _ScriptRunner,
refund_script: _ScriptRunner,
tokens: int,
scopes: tuple[BatchEnqueuedTokenScope, ...],
) -> BatchEnqueuedTokenOutcome:
started: Final = self._monotonic()
for index, scope in enumerate(scopes):
result = await self._run_reserve_script(
reserve_script,
refund_script,
tokens=tokens,
scope=scope,
already_reserved=scopes[:index],
)
if result[0] != 1:
await self._rollback_partial_reserve(refund_script, tokens=tokens, scopes=scopes[:index])
return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=result[1])
return BatchEnqueuedTokenReservation(
tokens=tokens, scopes=scopes, backend="redis", reserved_at_monotonic=started
)
async def _run_reserve_script(
self,
reserve_script: _ScriptRunner,
refund_script: _ScriptRunner,
tokens: int,
scope: BatchEnqueuedTokenScope,
already_reserved: tuple[BatchEnqueuedTokenScope, ...],
) -> tuple[int, int]:
try:
raw_result: Final = await reserve_script(
(self._counter_key(scope),),
(tokens, BATCH_ENQUEUED_TOKEN_TTL_SECONDS, scope.limit),
)
return _RESERVE_RESULT_ADAPTER.validate_python(raw_result)
except Exception:
await self._rollback_partial_reserve(refund_script, tokens=tokens, scopes=already_reserved)
raise
async def _rollback_partial_reserve(
self,
refund_script: _ScriptRunner,
tokens: int,
scopes: tuple[BatchEnqueuedTokenScope, ...],
) -> None:
try:
await self._refund_via_redis(refund_script, tokens=tokens, scopes=scopes)
except Exception as e: # noqa: BLE001 # best-effort rollback: the leak is TTL-bounded and only tightens the allowance
verbose_proxy_logger.warning(
"Rollback of partially reserved enqueued tokens failed; leaked increments expire with the TTL: %s",
str(e),
)
async def _refund_via_redis(
self,
refund_script: _ScriptRunner,
tokens: int,
scopes: tuple[BatchEnqueuedTokenScope, ...],
) -> None:
for scope in scopes:
await refund_script((self._counter_key(scope),), (tokens,))
async def _reserve_in_memory(
self,
tokens: int,
scopes: tuple[BatchEnqueuedTokenScope, ...],
span: "Span | None",
) -> BatchEnqueuedTokenOutcome:
started: Final = self._monotonic()
async with self._lock:
currents: Final = tuple([await self._get_local_counter(scope, span) for scope in scopes])
for scope, current in zip(scopes, currents):
if current + tokens > scope.limit:
return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=current)
for scope, current in zip(scopes, currents):
await self._set_local_counter(scope, current + tokens, span)
return BatchEnqueuedTokenReservation(
tokens=tokens, scopes=scopes, backend="memory", owner=self._owner_token, reserved_at_monotonic=started
)
async def refund(
self,
reservation: BatchEnqueuedTokenReservation,
litellm_parent_otel_span: "Span | None" = None,
) -> None:
if reservation.tokens <= 0 or not reservation.scopes:
return
if reservation.backend == "redis":
await self._refund_redis_reservation(reservation)
return
if reservation.owner != self._owner_token:
verbose_proxy_logger.warning(
"Skipping enqueued-token refund granted in another worker's memory; its counters expire with the TTL"
)
return
async with self._lock:
for scope in reservation.scopes:
current = await self._get_local_counter(scope, litellm_parent_otel_span)
remaining = current - reservation.tokens
if remaining <= 0:
self.internal_usage_cache.dual_cache.in_memory_cache.delete_cache(key=self._counter_key(scope))
else:
await self._set_local_counter(scope, remaining, litellm_parent_otel_span)
async def _refund_redis_reservation(self, reservation: BatchEnqueuedTokenReservation) -> None:
refund_script: Final = self._refund_script
if refund_script is None:
verbose_proxy_logger.warning(
"No Redis client for a Redis-granted enqueued-token refund; leaked increments expire with the TTL"
)
return
try:
await self._refund_via_redis(refund_script, tokens=reservation.tokens, scopes=reservation.scopes)
except Exception as e: # noqa: BLE001 # best-effort refund: the leak is TTL-bounded and only tightens the allowance
verbose_proxy_logger.warning(
"Redis enqueued-token refund failed; leaked increments expire with the TTL: %s", str(e)
)
async def save_reservation(
self,
batch_id: str,
reservation: BatchEnqueuedTokenReservation,
litellm_parent_otel_span: "Span | None" = None,
) -> None:
serialized: Final = _RESERVATION_ADAPTER.dump_json(reservation).decode("utf-8")
elapsed: Final = self._monotonic() - reservation.reserved_at_monotonic
ttl: Final = max(1, BATCH_ENQUEUED_TOKEN_TTL_SECONDS - math.ceil(elapsed))
if self._save_script is not None:
try:
await self._save_script(
(self._record_key(batch_id),),
(serialized, ttl),
)
except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record
verbose_proxy_logger.warning(
"Redis enqueued-token reservation save failed, falling back to in-memory: %s", str(e)
)
else:
return
await self.internal_usage_cache.async_set_cache(
key=self._record_key(batch_id),
value=serialized,
ttl=ttl,
litellm_parent_otel_span=litellm_parent_otel_span,
local_only=True,
)
async def pop_reservation(
self,
batch_id: str,
litellm_parent_otel_span: "Span | None" = None,
) -> BatchEnqueuedTokenReservation | None:
redis_raw: Final = await self._pop_redis_record(batch_id)
if redis_raw is not None and not redis_raw:
# The Redis pop tombstones popped records in place, so a hit on the empty
# tombstone means the batch was already refunded elsewhere; a local copy
# left behind by a save that raised after landing must not refund again.
await self._pop_local_record(batch_id, litellm_parent_otel_span)
return None
raw: Final = (
redis_raw if redis_raw is not None else await self._pop_local_record(batch_id, litellm_parent_otel_span)
)
if raw is None:
return None
try:
if isinstance(raw, (str, bytes)):
return _RESERVATION_ADAPTER.validate_json(raw)
return _RESERVATION_ADAPTER.validate_python(raw)
except ValidationError:
verbose_proxy_logger.warning("Discarding malformed enqueued-token reservation record for %s", batch_id)
return None
async def _pop_redis_record(self, batch_id: str) -> str | bytes | None:
pop_script: Final = self._pop_script
if pop_script is None:
return None
try:
return _POPPED_VALUE_ADAPTER.validate_python(
await pop_script((self._record_key(batch_id),), (BATCH_ENQUEUED_TOKEN_TTL_SECONDS,))
)
except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record
verbose_proxy_logger.warning(
"Redis enqueued-token reservation pop failed, falling back to in-memory: %s", str(e)
)
return None
async def _pop_local_record(self, batch_id: str, span: "Span | None") -> object:
async with self._lock:
stored = await self.internal_usage_cache.async_get_cache(
key=self._record_key(batch_id),
litellm_parent_otel_span=span,
local_only=True,
)
if stored is None:
return None
self.internal_usage_cache.dual_cache.in_memory_cache.delete_cache(key=self._record_key(batch_id))
return stored
async def _get_local_counter(self, scope: BatchEnqueuedTokenScope, span: "Span | None") -> int:
stored = await self.internal_usage_cache.async_get_cache(
key=self._counter_key(scope),
litellm_parent_otel_span=span,
local_only=True,
)
return _STORED_COUNTER_ADAPTER.validate_python(stored) or 0
async def _set_local_counter(self, scope: BatchEnqueuedTokenScope, value: int, span: "Span | None") -> None:
await self.internal_usage_cache.async_set_cache(
key=self._counter_key(scope),
value=value,
ttl=BATCH_ENQUEUED_TOKEN_TTL_SECONDS,
litellm_parent_otel_span=span,
local_only=True,
)

80
litellm/proxy/hooks/batch_rate_limiter.py
View file

@ -46,9 +46,16 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import (
ProxyRateLimitError,
map_v3_rate_limit_type,
)
from litellm.proxy.hooks.batch_enqueued_tokens import (
BatchEnqueuedTokenOverLimit,
BatchEnqueuedTokenReservation,
BatchEnqueuedTokenScope,
resolve_batch_enqueued_token_scopes,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROJECT_ITPM_DESCRIPTOR_KEY,
PROJECT_OTPM_DESCRIPTOR_KEY,
get_or_create_request_stash,
)
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
@ -291,6 +298,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
has_enqueued_scopes: bool = False,
) -> tuple[bool, list["RateLimitDescriptor"] | None]:
"""
Skip downloading batch input files when the operator disabled batch
@ -343,8 +351,10 @@ class _PROXY_BatchRateLimiter(CustomLogger):
user_api_key_dict=user_api_key_dict,
data=data,
)
if not self._has_applicable_batch_rate_limits(descriptors) and not self._project_has_any_io_token_limits(
user_api_key_dict
if (
not has_enqueued_scopes
and not self._has_applicable_batch_rate_limits(descriptors)
and not self._project_has_any_io_token_limits(user_api_key_dict)
):
verbose_proxy_logger.debug("Skipping batch input file processing: no rate limits configured")
return True, None
@ -511,6 +521,59 @@ class _PROXY_BatchRateLimiter(CustomLogger):
return file_id, fetch_kwargs
async def _reserve_batch_enqueued_tokens(
self,
user_api_key_dict: UserAPIKeyAuth,
data: Mapping[str, object],
batch_usage: BatchFileUsage,
scopes: tuple[BatchEnqueuedTokenScope, ...],
) -> None:
"""Reserve the batch's estimated tokens against the caller's enqueued-token allowance.
Runs instead of the per-minute counter charge when the key or team
opted in via ``batch_enqueued_token_limit`` metadata. The reservation
is stashed on the request so the v3 limiter's post-call hooks can
persist it (keyed by the provider batch id) and refund it when the
batch reaches a terminal state.
"""
outcome: Final = await self.parallel_request_limiter.batch_enqueued_token_store.reserve(
tokens=batch_usage.total_tokens,
scopes=scopes,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
match outcome:
case BatchEnqueuedTokenOverLimit():
self._raise_enqueued_limit_error(over_limit=outcome, data=data, batch_usage=batch_usage)
case BatchEnqueuedTokenReservation():
get_or_create_request_stash().batch_enqueued_reservation = outcome
def _raise_enqueued_limit_error(
self,
over_limit: BatchEnqueuedTokenOverLimit,
data: Mapping[str, object],
batch_usage: BatchFileUsage,
) -> NoReturn:
scope: Final = over_limit.scope
remaining: Final = max(0, scope.limit - over_limit.enqueued)
detail: Final = (
f"Batch enqueued token limit exceeded for {scope.key}: {scope.value}. "
f"Batch requires {batch_usage.total_tokens} tokens but only {remaining} enqueued tokens remaining "
f"out of {scope.limit} enqueued token limit. "
f"Tokens free up as running batches complete or are cancelled."
)
raw_model: Final = data.get("model")
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(
raw_model if isinstance(raw_model, str) else None
)
raise ProxyRateLimitError(
detail=detail,
headers=MappingProxyType({"rate_limit_type": "tokens"}),
category=RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT,
rate_limit_type=map_v3_rate_limit_type("tokens"),
model=resolved_model,
llm_provider=llm_provider,
)
def _raise_rate_limit_error(
self,
status: "RateLimitStatus",
@ -1039,8 +1102,9 @@ class _PROXY_BatchRateLimiter(CustomLogger):
verbose_proxy_logger.debug("No input_file_id in batch request, skipping rate limiting")
return data
enqueued_scopes: Final = resolve_batch_enqueued_token_scopes(user_api_key_dict)
should_skip, batch_rate_limit_descriptors = self._should_skip_batch_input_file_processing(
data=data, user_api_key_dict=user_api_key_dict
data=data, user_api_key_dict=user_api_key_dict, has_enqueued_scopes=bool(enqueued_scopes)
)
if should_skip:
return data
@ -1066,6 +1130,16 @@ class _PROXY_BatchRateLimiter(CustomLogger):
data["_batch_token_count"] = batch_usage.total_tokens
data["_batch_request_count"] = batch_usage.request_count
if enqueued_scopes:
await self._reserve_batch_enqueued_tokens(
user_api_key_dict=user_api_key_dict,
data=data,
batch_usage=batch_usage,
scopes=enqueued_scopes,
)
verbose_proxy_logger.debug("Batch enqueued-token reservation succeeded")
return data
# Directly increment counters by batch amounts (check happens atomically)
# This will raise HTTPException if limits are exceeded
await self._check_and_increment_batch_counters(

42
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -44,6 +44,13 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import (
ProxyRateLimitError,
map_v3_rate_limit_type,
)
from litellm.proxy.hooks.batch_enqueued_tokens import (
BATCH_ENQUEUED_REFUND_STATUSES,
BatchEnqueuedTokenReservation,
BatchEnqueuedTokenStore,
batch_response_view,
canonical_provider_batch_id,
)
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage
@ -515,6 +522,7 @@ class RequestRateLimiterStash:
otpm_reserved_window_identities: frozenset[tuple[str, str, Literal["redis", "local"]]] = field(
default_factory=frozenset
)
batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None
reservation_released: bool = False
@ -619,6 +627,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Batch rate limiter (lazy loaded)
self._batch_rate_limiter: CallTypeRateLimiter | None = None
self.batch_enqueued_token_store = BatchEnqueuedTokenStore(internal_usage_cache=internal_usage_cache)
# Serializes multi-phase check+increment sequences (batch + dynamic
# limiters) within this process to close the TOCTOU window between
@ -4673,6 +4682,32 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
except Exception as e:
verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e)
try:
await self._handle_batch_enqueued_post_call(user_api_key_dict=user_api_key_dict, response=response)
except Exception as e: # noqa: BLE001 # post-call batch accounting must never fail the response
verbose_proxy_logger.exception("Error in batch enqueued-token post-call hook: %s", e)
async def _handle_batch_enqueued_post_call(self, user_api_key_dict: UserAPIKeyAuth, response: object) -> None:
view: Final = batch_response_view(response)
if view is None:
return
span: Final = user_api_key_dict.parent_otel_span
stash: Final = get_request_stash()
if stash is not None and stash.batch_enqueued_reservation is not None:
await self.batch_enqueued_token_store.save_reservation(
batch_id=canonical_provider_batch_id(view.id),
reservation=stash.batch_enqueued_reservation,
litellm_parent_otel_span=span,
)
stash.batch_enqueued_reservation = None
if view.status.lower() in BATCH_ENQUEUED_REFUND_STATUSES:
popped: Final = await self.batch_enqueued_token_store.pop_reservation(
batch_id=canonical_provider_batch_id(view.id),
litellm_parent_otel_span=span,
)
if popped is not None:
await self.batch_enqueued_token_store.refund(reservation=popped, litellm_parent_otel_span=span)
async def async_post_call_failure_hook(
self,
request_data: dict,
@ -4706,6 +4741,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
stash.parallel_slot = None
if stash.batch_enqueued_reservation is not None:
await self.batch_enqueued_token_store.refund(
reservation=stash.batch_enqueued_reservation,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
stash.batch_enqueued_reservation = None
if stash.reservation_released:
return
reserved_tokens: Final = stash.reserved_tokens

27
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -57,6 +57,7 @@ from litellm.proxy.auth.auth_checks import (
)
from litellm.proxy.auth.auth_utils import (
abbreviate_api_key,
enforce_batch_enqueued_token_limit_is_admin_only,
enforce_output_token_estimates_are_admin_only,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -901,6 +902,12 @@ async def _common_key_generation_helper(
user_api_key_dict=user_api_key_dict,
entity="key",
)
enforce_batch_enqueued_token_limit_is_admin_only(
data=data,
existing_metadata=None,
user_api_key_dict=user_api_key_dict,
entity="key",
)
if data.metadata is not None and data.metadata.get("service_account_id") is not None and data.team_id is None:
await validate_team_id_used_in_service_account_request(
@ -2302,6 +2309,14 @@ async def _process_single_key_update(
prisma_client=prisma_client,
)
_existing_row_metadata: Final = getattr(existing_key_row, "metadata", None)
enforce_batch_enqueued_token_limit_is_admin_only(
data=update_key_request,
existing_metadata=_existing_row_metadata if isinstance(_existing_row_metadata, dict) else None,
user_api_key_dict=user_api_key_dict,
entity="key",
)
# Check team member permissions
if prisma_client is not None:
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
@ -2564,6 +2579,12 @@ async def _validate_update_key_data(
user_api_key_dict=user_api_key_dict,
entity="key",
)
enforce_batch_enqueued_token_limit_is_admin_only(
data=data,
existing_metadata=_existing_metadata if isinstance(_existing_metadata, dict) else None,
user_api_key_dict=user_api_key_dict,
entity="key",
)
# Personal-key bypass: the caller both created the key AND still owns it
# (user_id == caller). Checking only created_by would let a demoted admin
@ -4760,6 +4781,12 @@ async def _execute_virtual_key_regeneration(
user_api_key_dict=user_api_key_dict,
entity="key",
)
enforce_batch_enqueued_token_limit_is_admin_only(
data=data,
existing_metadata=_existing_key_metadata if isinstance(_existing_key_metadata, dict) else None,
user_api_key_dict=user_api_key_dict,
entity="key",
)
new_token: Final = await get_new_token(data=data)
new_token_hash: Final = hash_token(new_token)

17
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -85,7 +85,10 @@ from litellm.proxy.auth.auth_checks import (
get_team_object,
get_user_object,
)
from litellm.proxy.auth.auth_utils import enforce_output_token_estimates_are_admin_only
from litellm.proxy.auth.auth_utils import (
enforce_batch_enqueued_token_limit_is_admin_only,
enforce_output_token_estimates_are_admin_only,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch
@ -1304,6 +1307,12 @@ async def new_team(
user_api_key_dict=user_api_key_dict,
entity="team",
)
enforce_batch_enqueued_token_limit_is_admin_only(
data=data,
existing_metadata=None,
user_api_key_dict=user_api_key_dict,
entity="team",
)
# Check if license is over limit
total_teams: Final = await _team_db(prisma_client).count()
@ -2008,6 +2017,12 @@ async def update_team(
user_api_key_dict=user_api_key_dict,
entity="team",
)
enforce_batch_enqueued_token_limit_is_admin_only(
data=data,
existing_metadata=_existing_team_metadata if isinstance(_existing_team_metadata, dict) else None,
user_api_key_dict=user_api_key_dict,
entity="team",
)
_check_passthrough_routes_caller_permission(data, user_api_key_dict, entity="team")

146
tests/e2e/batches/test_batches_e2e.py
View file

@ -17,6 +17,7 @@ from __future__ import annotations
import json
import os
import re
import time
from datetime import datetime, timedelta, timezone
from typing import Callable
@ -57,7 +58,7 @@ from e2e_http import (
unwrap,
)
from lifecycle import ResourceManager
from models import KeyGenerateBody, LiteLLMParamsBody, SpendLogRow
from models import KeyGenerateBody, KeyMetadata, LiteLLMParamsBody, SpendLogRow
pytestmark = pytest.mark.e2e
@ -685,6 +686,149 @@ class TestBatchRateLimitErrorMapping:
)
BATCH_ENQUEUED_HEADROOM_TOKENS = 100_000
_BATCH_REQUIRES_TOKENS = re.compile(r"Batch requires (\d+) tokens")
class TestBatchEnqueuedTokenLimit:
"""Opt-in enqueued-token allowance governs batch submission instead of RPM/TPM.
A key whose metadata carries batch_enqueued_token_limit reserves the batch's
token estimate against that allowance at create time: per-minute limits no
longer gate batch submission, exhausting the allowance rejects the create
before it reaches the provider, and cancelling a running batch refunds its
reservation so blocked submissions go through again (LIT-5273).
"""
def _upload_batch_file(
self, client: BatchClient, resources: ResourceManager, key: str
) -> FileObject:
file = unwrap(
client.upload_file(
content=_multi_request_jsonl("gpt-4o-mini", BATCH_RL_REQUEST_LINES),
form=FileUploadForm(purpose="batch"),
model=OPENAI_BATCH_MODEL,
key=key,
)
)
resources.defer(quietly(lambda: client.delete_file(file.id, key=key)))
return file
def _generate_enqueued_key(
self,
client: BatchClient,
resources: ResourceManager,
*,
limit: int,
marker: str,
rpm_limit: int | None = None,
) -> str:
key = client.proxy.generate_key(
KeyGenerateBody(
models=[],
rpm_limit=rpm_limit,
user_id=f"e2e-batch-enq-{marker}-{unique_marker()}",
metadata=KeyMetadata(batch_enqueued_token_limit=limit),
)
)
resources.defer(lambda: client.proxy.delete_key(key))
return key
@pytest.mark.covers(
"quota_management.ratelimit.batch_enqueued_tokens.accepts_over_rpm",
exercised_on=["batches"],
)
def test_enqueued_allowance_accepts_batch_over_key_rpm(
self, client: BatchClient, resources: ResourceManager, batch_deployments: None
) -> None:
key = self._generate_enqueued_key(
client,
resources,
limit=BATCH_ENQUEUED_HEADROOM_TOKENS,
marker="rpm",
rpm_limit=BATCH_RL_RPM_LIMIT,
)
file = self._upload_batch_file(client, resources, key)
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
assert created.status_code != 429, (
f"enqueued-token allowance must govern batch submission instead of the "
f"key RPM ({BATCH_RL_RPM_LIMIT} < {BATCH_RL_REQUEST_LINES} rows); "
f"got 429: {created.body[:400]}"
)
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key)))
@pytest.mark.covers(
"quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted",
exercised_on=["batches"],
)
@pytest.mark.covers(
"quota_management.ratelimit.batch_enqueued_tokens.refunds_on_cancel",
exercised_on=["batches"],
)
def test_exhausted_allowance_blocks_until_cancel_refunds(
self, client: BatchClient, resources: ResourceManager, batch_deployments: None
) -> None:
sizing_key = self._generate_enqueued_key(
client, resources, limit=1, marker="size"
)
sizing_file = self._upload_batch_file(client, resources, sizing_key)
sized = client.create_batch(
body=BatchCreateBody(input_file_id=sizing_file.id), key=sizing_key
)
assert sized.status_code == 429, (
f"a 1-token allowance must reject any batch before it reaches the "
f"provider, got {sized.status_code}: {sized.body[:400]}"
)
assert "batch enqueued token limit exceeded" in sized.body.lower(), (
f"429 body must name the enqueued token limit, got: {sized.body[:400]}"
)
requires = _BATCH_REQUIRES_TOKENS.search(sized.body)
assert requires is not None, (
f"429 body must report the batch token requirement so callers can size "
f"allowances, got: {sized.body[:400]}"
)
batch_tokens = int(requires.group(1))
assert batch_tokens > 1
key = self._generate_enqueued_key(
client, resources, limit=batch_tokens + batch_tokens // 2, marker="refund"
)
file = self._upload_batch_file(client, resources, key)
first = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
require_successful_call(first)
first_batch = BatchObject.model_validate_json(first.body)
resources.defer(quietly(lambda: client.cancel_batch(first_batch.id, key=key)))
blocked = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
assert blocked.status_code == 429, (
f"second batch must not fit the remaining allowance while the first is "
f"enqueued, got {blocked.status_code}: {blocked.body[:400]}"
)
assert "batch enqueued token limit exceeded" in blocked.body.lower(), (
f"429 body must name the enqueued token limit, got: {blocked.body[:400]}"
)
cancelled = cancel_batch(client, first_batch.id, key=key, provider=None)
assert cancelled.status in {"cancelling", "cancelled"}, (
f"cancel must reach a cancel state for the refund to fire, "
f"got {cancelled.status}"
)
retried = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
assert retried.status_code != 429, (
f"cancelling the first batch must refund its reservation so the retry "
f"fits the allowance, got 429: {retried.body[:400]}"
)
require_successful_call(retried)
retry_batch = BatchObject.model_validate_json(retried.body)
resources.defer(quietly(lambda: client.cancel_batch(retry_batch.id, key=key)))
ASSUME_ROLE_RAW_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"

3
tests/e2e/coverage_registry/quota_management.yaml
View file

@ -2,6 +2,9 @@
# litellm/proxy/hooks/ + litellm/proxy/auth/auth_checks.py + litellm/proxy/spend_tracking/.
- {id: quota_management.ratelimit.rpm.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: rpm, assertions: [blocks_over_limit], exercised_on: [chat_completions, messages], source: "parallel_request_limiter_v3.py", rationale: "v3 limiter enforces RPM per key/team/model; 429 on breach"}
- {id: quota_management.ratelimit.batch_rpm.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: batch_rpm, assertions: [blocks_over_limit], exercised_on: [batches], source: "batch_rate_limiter.py", rationale: "Batch create that exceeds key RPM returns mapped 429 with retry-after"}
- {id: quota_management.ratelimit.batch_enqueued_tokens.accepts_over_rpm, module: quota_management, tier: P0, behavior: ratelimit, variant: batch_enqueued_tokens, assertions: [accepts_over_rpm], exercised_on: [batches], source: "batch_rate_limiter.py + batch_enqueued_tokens.py", rationale: "Key with an enqueued-token allowance submits a batch whose row count exceeds its RPM and the create is accepted"}
- {id: quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted, module: quota_management, tier: P0, behavior: ratelimit, variant: batch_enqueued_tokens, assertions: [blocks_when_exhausted], exercised_on: [batches], source: "batch_rate_limiter.py + batch_enqueued_tokens.py", rationale: "Batch create is rejected with a 429 naming the enqueued token limit once the allowance cannot fit the file"}
- {id: quota_management.ratelimit.batch_enqueued_tokens.refunds_on_cancel, module: quota_management, tier: P0, behavior: ratelimit, variant: batch_enqueued_tokens, assertions: [refunds_on_cancel], exercised_on: [batches], source: "batch_rate_limiter.py + batch_enqueued_tokens.py", rationale: "Cancelling a running batch returns its reserved tokens so a previously blocked submission succeeds"}
- {id: quota_management.ratelimit.tpm.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: tpm, assertions: [blocks_over_limit], exercised_on: [chat_completions, messages], source: "parallel_request_limiter_v3.py", rationale: "v3 limiter enforces TPM per key/team/model; 429 on breach"}
- {id: quota_management.ratelimit.tpm.excludes_cached_tokens, module: quota_management, tier: P0, behavior: ratelimit, variant: tpm, assertions: [excludes_cached_tokens], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py:_get_total_tokens_from_usage", rationale: "Cached prompt tokens must not count toward TPM (LIT-1930)"}
- {id: quota_management.ratelimit.redis_backed.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: redis_backed, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py", rationale: "With Redis configured, RPM still enforces 429 across the shared limiter path customers run multi-replica"}

1
tests/e2e/models.py
View file

@ -46,6 +46,7 @@ class KeyLoggingCallback(BaseModel):
class KeyMetadata(BaseModel):
logging: list[KeyLoggingCallback] | None = None
priority: str | None = None
batch_enqueued_token_limit: int | None = None
class ObjectPermission(BaseModel):

441
tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py Normal file
View file

@ -0,0 +1,441 @@
"""
LIT-5273: enqueued-token accounting for batch submissions.
Covers the ``BatchEnqueuedTokenStore`` (reserve / refund / reservation
records), the metadata-driven scope resolution, and the batch-id and
response-shape helpers the v3 limiter's post-call hooks rely on.
"""
import base64
import socket
import uuid
from collections.abc import Mapping, Sequence
from types import MappingProxyType, SimpleNamespace
from typing import Final
import pytest
from litellm.caching.caching import DualCache
from litellm.constants import BATCH_ENQUEUED_TOKEN_TTL_SECONDS
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.batch_enqueued_tokens import (
BatchEnqueuedTokenOverLimit,
BatchEnqueuedTokenReservation,
BatchEnqueuedTokenScope,
BatchEnqueuedTokenStore,
batch_response_view,
canonical_provider_batch_id,
resolve_batch_enqueued_token_scopes,
)
from litellm.proxy.utils import InternalUsageCache
def _in_memory_store() -> BatchEnqueuedTokenStore:
return BatchEnqueuedTokenStore(internal_usage_cache=InternalUsageCache(DualCache(default_in_memory_ttl=60)))
def _scope(limit: int, key: str = "api_key") -> BatchEnqueuedTokenScope:
return BatchEnqueuedTokenScope(key=key, value=f"{key}-{uuid.uuid4().hex}", limit=limit)
def test_scope_resolution_reads_key_and_team_metadata():
user = UserAPIKeyAuth(
api_key="hashed-key",
metadata={"batch_enqueued_token_limit": 100},
team_id="team-1",
team_metadata={"batch_enqueued_token_limit": "150"},
)
scopes = resolve_batch_enqueued_token_scopes(user)
assert scopes == (
BatchEnqueuedTokenScope(key="api_key", value="hashed-key", limit=100),
BatchEnqueuedTokenScope(key="team", value="team-1", limit=150),
)
def test_scope_resolution_returns_empty_without_opt_in():
assert resolve_batch_enqueued_token_scopes(UserAPIKeyAuth(api_key="k")) == ()
assert resolve_batch_enqueued_token_scopes(UserAPIKeyAuth(api_key="k", metadata={}, team_metadata=None)) == ()
@pytest.mark.parametrize("bad_value", ["not-a-number", 0, -5, None, [1000]])
def test_scope_resolution_ignores_invalid_limits(bad_value):
user = UserAPIKeyAuth(api_key="k", metadata={"batch_enqueued_token_limit": bad_value})
assert resolve_batch_enqueued_token_scopes(user) == ()
def test_scope_resolution_skips_team_scope_without_team_id():
user = UserAPIKeyAuth(api_key="k", team_metadata={"batch_enqueued_token_limit": 100})
assert resolve_batch_enqueued_token_scopes(user) == ()
@pytest.mark.asyncio
async def test_reserve_rejects_once_allowance_is_exhausted():
store = _in_memory_store()
scope = _scope(limit=100)
first = await store.reserve(tokens=80, scopes=(scope,))
assert isinstance(first, BatchEnqueuedTokenReservation)
second = await store.reserve(tokens=30, scopes=(scope,))
assert second == BatchEnqueuedTokenOverLimit(scope=scope, enqueued=80)
third = await store.reserve(tokens=20, scopes=(scope,))
assert isinstance(third, BatchEnqueuedTokenReservation)
@pytest.mark.asyncio
async def test_reserve_is_all_or_nothing_across_scopes():
store = _in_memory_store()
key_scope = _scope(limit=100, key="api_key")
team_scope = _scope(limit=50, key="team")
over = await store.reserve(tokens=60, scopes=(key_scope, team_scope))
assert over == BatchEnqueuedTokenOverLimit(scope=team_scope, enqueued=0)
exact_fit = await store.reserve(tokens=50, scopes=(key_scope, team_scope))
assert isinstance(exact_fit, BatchEnqueuedTokenReservation)
@pytest.mark.asyncio
async def test_refund_restores_allowance_and_never_goes_negative():
store = _in_memory_store()
scope = _scope(limit=100)
reservation = await store.reserve(tokens=30, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
await store.refund(reservation)
await store.refund(reservation)
refill = await store.reserve(tokens=100, scopes=(scope,))
assert isinstance(refill, BatchEnqueuedTokenReservation)
assert isinstance(await store.reserve(tokens=1, scopes=(scope,)), BatchEnqueuedTokenOverLimit)
@pytest.mark.asyncio
async def test_reservation_record_roundtrip_pops_exactly_once():
store = _in_memory_store()
scope = _scope(limit=100)
reservation = await store.reserve(tokens=40, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
await store.save_reservation("batch_abc", reservation)
assert await store.pop_reservation("batch_abc") == reservation
assert await store.pop_reservation("batch_abc") is None
assert await store.pop_reservation("batch_never_saved") is None
@pytest.mark.asyncio
async def test_zero_token_reserve_charges_nothing():
store = _in_memory_store()
scope = _scope(limit=100)
empty = await store.reserve(tokens=0, scopes=(scope,))
assert empty == BatchEnqueuedTokenReservation(tokens=0, scopes=(scope,))
full = await store.reserve(tokens=100, scopes=(scope,))
assert isinstance(full, BatchEnqueuedTokenReservation)
class _SingleKeyRedisFake:
"""Emulates the Redis script path one single-key call at a time, recording every call."""
def __init__(
self,
fail_reserve_keys: frozenset[str] = frozenset(),
fail_refund_keys: frozenset[str] = frozenset(),
fail_save_keys: frozenset[str] = frozenset(),
raise_after_landing_save_keys: frozenset[str] = frozenset(),
) -> None:
self.script_calls: tuple[tuple[str, tuple[str, ...]], ...] = ()
self.save_ttls: tuple[int, ...] = ()
self.counters: Mapping[str, int] = MappingProxyType({})
self.records: Mapping[str, str] = MappingProxyType({})
self.fail_reserve_keys = fail_reserve_keys
self.fail_refund_keys = fail_refund_keys
self.fail_save_keys = fail_save_keys
self.raise_after_landing_save_keys = raise_after_landing_save_keys
def async_register_script(self, script: str):
kind: Final = (
"reserve"
if "INCRBY" in script
else "refund" if "DECRBY" in script else "pop" if "GET" in script else "save"
)
async def run(keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> object:
self.script_calls = (*self.script_calls, (kind, tuple(keys)))
return self._run(kind, tuple(keys), tuple(args))
return run
def _run(self, kind: str, keys: tuple[str, ...], args: tuple[str | bytes | int | float, ...]) -> object:
if kind == "reserve":
if keys[0] in self.fail_reserve_keys:
raise ConnectionError(f"simulated redis failure for {keys[0]}")
amount, limit = int(args[0]), int(args[2])
current: Final = self.counters.get(keys[0], 0)
if current + amount > limit:
return (0, current)
self.counters = MappingProxyType({**self.counters, keys[0]: current + amount})
return (1, current + amount)
if kind == "refund":
if keys[0] in self.fail_refund_keys:
raise ConnectionError(f"simulated redis failure for {keys[0]}")
remaining: Final = self.counters.get(keys[0], 0) - int(args[0])
self.counters = MappingProxyType(
{key: value for key, value in self.counters.items() if key != keys[0]}
if remaining <= 0
else {**self.counters, keys[0]: remaining}
)
return 1
if kind == "save":
if keys[0] in self.fail_save_keys:
raise ConnectionError(f"simulated redis failure for {keys[0]}")
self.records = MappingProxyType({**self.records, keys[0]: str(args[0])})
self.save_ttls = (*self.save_ttls, int(args[1]))
if keys[0] in self.raise_after_landing_save_keys:
raise TimeoutError(f"simulated redis timeout after landing for {keys[0]}")
return 1
if kind == "pop":
popped: Final = self.records.get(keys[0])
if popped:
self.records = MappingProxyType({**self.records, keys[0]: ""})
return popped
raise AssertionError(f"unexpected {kind} script call for keys {keys}")
@pytest.mark.asyncio
async def test_redis_reserve_issues_single_key_calls_and_rolls_back_on_over_limit():
fake = _SingleKeyRedisFake()
store = BatchEnqueuedTokenStore(
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
)
key_scope = _scope(limit=100, key="api_key")
team_scope = _scope(limit=50, key="team")
over = await store.reserve(tokens=60, scopes=(key_scope, team_scope))
assert over == BatchEnqueuedTokenOverLimit(scope=team_scope, enqueued=0)
assert tuple(kind for kind, _ in fake.script_calls) == ("reserve", "reserve", "refund")
assert not fake.counters
fits = await store.reserve(tokens=50, scopes=(key_scope, team_scope))
assert isinstance(fits, BatchEnqueuedTokenReservation)
await store.refund(fits)
assert not fake.counters
assert all(len(keys) == 1 for _, keys in fake.script_calls)
@pytest.mark.asyncio
async def test_partial_redis_reserve_failure_rolls_back_and_grants_in_memory():
key_scope = _scope(limit=100, key="api_key")
team_scope = _scope(limit=50, key="team")
fake = _SingleKeyRedisFake(fail_reserve_keys=frozenset({f"batch_enqueued_tokens:team:{team_scope.value}"}))
store = BatchEnqueuedTokenStore(
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
)
outcome = await store.reserve(tokens=10, scopes=(key_scope, team_scope))
assert isinstance(outcome, BatchEnqueuedTokenReservation)
assert outcome.backend == "memory"
assert tuple(kind for kind, _ in fake.script_calls) == ("reserve", "reserve", "refund")
assert not fake.counters
await store.refund(outcome)
assert tuple(kind for kind, _ in fake.script_calls) == ("reserve", "reserve", "refund")
refilled = await store.reserve(tokens=50, scopes=(team_scope,))
assert isinstance(refilled, BatchEnqueuedTokenReservation)
assert refilled.backend == "memory"
@pytest.mark.asyncio
async def test_over_limit_verdict_survives_a_failing_rollback():
key_scope = _scope(limit=100, key="api_key")
team_scope = _scope(limit=5, key="team")
key_counter: Final = f"batch_enqueued_tokens:api_key:{key_scope.value}"
fake = _SingleKeyRedisFake(fail_refund_keys=frozenset({key_counter}))
store = BatchEnqueuedTokenStore(
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
)
outcome = await store.reserve(tokens=10, scopes=(key_scope, team_scope))
assert outcome == BatchEnqueuedTokenOverLimit(scope=team_scope, enqueued=0)
assert fake.counters == {key_counter: 10}
@pytest.mark.asyncio
async def test_pop_falls_back_to_local_record_when_redis_pop_finds_nothing():
scope = _scope(limit=100)
record_key: Final = "batch_enqueued_token_reservation:batch_local_record"
fake = _SingleKeyRedisFake(fail_save_keys=frozenset({record_key}))
store = BatchEnqueuedTokenStore(
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
)
reservation = await store.reserve(tokens=60, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
await store.save_reservation("batch_local_record", reservation)
assert not fake.records
popped = await store.pop_reservation("batch_local_record")
assert popped == reservation
await store.refund(popped)
assert not fake.counters
assert await store.pop_reservation("batch_local_record") is None
@pytest.mark.asyncio
async def test_local_ghost_left_by_landed_save_never_refunds_twice():
scope = _scope(limit=100)
record_key: Final = "batch_enqueued_token_reservation:batch_ghost"
fake = _SingleKeyRedisFake(raise_after_landing_save_keys=frozenset({record_key}))
store = BatchEnqueuedTokenStore(
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
)
reservation = await store.reserve(tokens=60, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
await store.save_reservation("batch_ghost", reservation)
assert fake.records[record_key]
first = await store.pop_reservation("batch_ghost")
assert first == reservation
await store.refund(first)
assert not fake.counters
assert fake.records[record_key] == ""
assert await store.pop_reservation("batch_ghost") is None
assert (
await store.internal_usage_cache.async_get_cache(
key=record_key, litellm_parent_otel_span=None, local_only=True
)
is None
)
assert await store.pop_reservation("batch_ghost") is None
assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenReservation)
assert fake.counters[f"batch_enqueued_tokens:{scope.key}:{scope.value}"] == 100
@pytest.mark.asyncio
async def test_record_ttl_shrinks_by_elapsed_time_so_stale_records_never_outlive_their_counters():
scope = _scope(limit=100)
fake = _SingleKeyRedisFake()
ticks = iter((1_000.0, 1_030.5))
store = BatchEnqueuedTokenStore(
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60)),
monotonic=lambda: next(ticks),
)
reservation = await store.reserve(tokens=60, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
assert reservation.reserved_at_monotonic == 1_000.0
await store.save_reservation("batch_ttl_clamp", reservation)
assert fake.save_ttls == (BATCH_ENQUEUED_TOKEN_TTL_SECONDS - 31,)
assert await store.pop_reservation("batch_ttl_clamp") == reservation
@pytest.mark.asyncio
async def test_memory_refund_skips_reservations_granted_by_another_worker():
store = _in_memory_store()
scope = _scope(limit=100)
reservation = await store.reserve(tokens=60, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
assert reservation.backend == "memory"
assert reservation.owner
foreign: Final = BatchEnqueuedTokenReservation(
tokens=60, scopes=reservation.scopes, backend="memory", owner="another-worker"
)
await store.refund(foreign)
assert await store.reserve(tokens=50, scopes=(scope,)) == BatchEnqueuedTokenOverLimit(scope=scope, enqueued=60)
await store.refund(reservation)
assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenReservation)
@pytest.mark.asyncio
async def test_failed_redis_refund_leaves_local_counters_untouched():
scope = _scope(limit=100)
counter_key: Final = f"batch_enqueued_tokens:api_key:{scope.value}"
fake = _SingleKeyRedisFake(fail_refund_keys=frozenset({counter_key}))
store = BatchEnqueuedTokenStore(
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
)
reservation = await store.reserve(tokens=60, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
assert reservation.backend == "redis"
store.internal_usage_cache.dual_cache.in_memory_cache.set_cache(key=counter_key, value=45)
await store.refund(reservation)
assert store.internal_usage_cache.dual_cache.in_memory_cache.get_cache(key=counter_key) == 45
assert fake.counters == {counter_key: 60}
@pytest.mark.asyncio
async def test_pop_reservation_defaults_legacy_records_to_redis_backend():
store = _in_memory_store()
legacy = '{"tokens": 5, "scopes": [{"key": "api_key", "value": "k", "limit": 10}]}'
store.internal_usage_cache.dual_cache.in_memory_cache.set_cache(
key="batch_enqueued_token_reservation:batch_legacy", value=legacy
)
popped = await store.pop_reservation("batch_legacy")
assert popped == BatchEnqueuedTokenReservation(
tokens=5, scopes=(BatchEnqueuedTokenScope(key="api_key", value="k", limit=10),), backend="redis"
)
def test_canonical_provider_batch_id_passes_raw_ids_through():
assert canonical_provider_batch_id("batch_abc123") == "batch_abc123"
def test_canonical_provider_batch_id_decodes_unified_batch_ids():
unified = "litellm_proxy;model_id:m-1;llm_batch_id:batch_prov_9;llm_output_file_id:file-9"
encoded = base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")
assert canonical_provider_batch_id(encoded) == "batch_prov_9"
def test_canonical_provider_batch_id_decodes_model_embedded_ids():
from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
encoded = encode_file_id_with_model(file_id="batch_prov_7", model="my-alias", id_type="batch")
assert canonical_provider_batch_id(encoded) == "batch_prov_7"
def test_batch_response_view_accepts_batch_objects_only():
batch = SimpleNamespace(id="batch_1", status="completed", object="batch")
view = batch_response_view(batch)
assert view is not None and view.id == "batch_1" and view.status == "completed"
assert batch_response_view({"id": "chatcmpl-1", "object": "chat.completion"}) is None
assert batch_response_view(None) is None
assert batch_response_view("batch_1") is None
def _local_redis_port() -> int | None:
for port in (6379,):
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.settimeout(0.2)
if sock.connect_ex(("127.0.0.1", port)) == 0:
return port
return None
@pytest.mark.asyncio
@pytest.mark.skipif(_local_redis_port() is None, reason="requires a local Redis on 6379 for the Lua script path")
async def test_redis_lua_path_full_lifecycle():
from litellm.caching.redis_cache import RedisCache
port = _local_redis_port()
redis_cache = RedisCache(host="127.0.0.1", port=port)
store = BatchEnqueuedTokenStore(
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis_cache, default_in_memory_ttl=60))
)
key_scope = _scope(limit=100, key="api_key")
team_scope = _scope(limit=50, key="team")
over = await store.reserve(tokens=60, scopes=(key_scope, team_scope))
assert over == BatchEnqueuedTokenOverLimit(scope=team_scope, enqueued=0)
reservation = await store.reserve(tokens=50, scopes=(key_scope, team_scope))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
assert isinstance(await store.reserve(tokens=1, scopes=(key_scope, team_scope)), BatchEnqueuedTokenOverLimit)
batch_id = f"batch_{uuid.uuid4().hex}"
await store.save_reservation(batch_id, reservation)
popped = await store.pop_reservation(batch_id)
assert popped == reservation
assert await store.pop_reservation(batch_id) is None
await store.refund(popped)
refill = await store.reserve(tokens=50, scopes=(key_scope, team_scope))
assert isinstance(refill, BatchEnqueuedTokenReservation)
await store.refund(refill)

192
tests/test_litellm/proxy/hooks/test_batch_file_validation.py
View file

@ -2094,3 +2094,195 @@ def test_estimate_entry_output_tokens_multiplies_candidate_count(body_extra, exp
}
assert rate_limiter._estimate_entry_output_tokens(entry, None) == expected
# ---------------------------------------------------------------------------
# LIT-5273: enqueued-token limits govern batch submission when opted in
# ---------------------------------------------------------------------------
def _enqueued_rate_limiter():
from litellm import DualCache
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.utils import InternalUsageCache
local_cache = DualCache(default_in_memory_ttl=60)
internal_usage_cache = InternalUsageCache(local_cache)
parallel_request_limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache)
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=internal_usage_cache,
parallel_request_limiter=parallel_request_limiter,
)
return rate_limiter, local_cache
_ENQUEUED_BATCH_FILE_CONTENT = (
b'{"body": {"model": "gpt-4o-mini", "max_tokens": 40, "messages": [{"role": "user", "content": "hi"}]}}\n'
b'{"body": {"model": "gpt-4o-mini", "max_tokens": 40, "messages": [{"role": "user", "content": "hi"}]}}\n'
b'{"body": {"model": "gpt-4o-mini", "max_tokens": 40, "messages": [{"role": "user", "content": "hi"}]}}\n'
)
def _enqueued_batch_patches():
mock_content = MagicMock()
mock_content.content = _ENQUEUED_BATCH_FILE_CONTENT
afile_content_mock = AsyncMock(return_value=mock_content)
return afile_content_mock, (
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch(
"litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model",
return_value={"custom_llm_provider": "openai"},
),
)
@pytest.mark.asyncio
async def test_enqueued_limit_accepts_batch_over_per_minute_limits():
"""The headline LIT-5273 behavior: a key that opted into an enqueued-token
allowance submits a batch whose row count and token count both exceed its
per-minute RPM/TPM limits, and the batch is accepted (repeatedly) because
only the enqueued allowance governs. Without the opt-in the same key is
rejected on RPM before the batch reaches the provider."""
from litellm.proxy.hooks.parallel_request_limiter_v3 import get_request_stash
rate_limiter, local_cache = _enqueued_rate_limiter()
afile_content_mock, patches = _enqueued_batch_patches()
legacy_user = UserAPIKeyAuth(api_key="sk-legacy-rpm", models=["*"], rpm_limit=1, tpm_limit=10)
opted_in_user = UserAPIKeyAuth(
api_key="sk-enqueued-rpm",
models=["*"],
rpm_limit=1,
tpm_limit=10,
metadata={"batch_enqueued_token_limit": 100000},
)
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
with pytest.raises(HTTPException) as legacy_exc:
await rate_limiter.async_pre_call_hook(
user_api_key_dict=legacy_user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
assert legacy_exc.value.status_code == 429
first_data = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
result = await rate_limiter.async_pre_call_hook(
user_api_key_dict=opted_in_user,
cache=local_cache,
data=first_data,
call_type="acreate_batch",
)
assert result is first_data
stash = get_request_stash()
assert stash is not None and stash.batch_enqueued_reservation is not None
assert stash.batch_enqueued_reservation.tokens == first_data["_batch_token_count"] > 0
second = await rate_limiter.async_pre_call_hook(
user_api_key_dict=opted_in_user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
assert second is not None
@pytest.mark.asyncio
async def test_enqueued_limit_rejects_when_allowance_is_exhausted():
"""Submissions are rejected pre-provider once the enqueued allowance can't
fit the batch, even for a key with no per-minute limits at all (which
previously skipped batch rate limiting entirely)."""
rate_limiter, local_cache = _enqueued_rate_limiter()
afile_content_mock, patches = _enqueued_batch_patches()
sizing_user = UserAPIKeyAuth(
api_key="sk-enqueued-sizing", models=["*"], metadata={"batch_enqueued_token_limit": 1000000}
)
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
sizing_data = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
await rate_limiter.async_pre_call_hook(
user_api_key_dict=sizing_user,
cache=local_cache,
data=sizing_data,
call_type="acreate_batch",
)
batch_tokens = sizing_data["_batch_token_count"]
assert batch_tokens > 0
capped_user = UserAPIKeyAuth(
api_key="sk-enqueued-capped",
models=["*"],
metadata={"batch_enqueued_token_limit": batch_tokens + batch_tokens // 2},
)
await rate_limiter.async_pre_call_hook(
user_api_key_dict=capped_user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
with pytest.raises(HTTPException) as exc:
await rate_limiter.async_pre_call_hook(
user_api_key_dict=capped_user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
assert exc.value.status_code == 429
assert "Batch enqueued token limit exceeded for api_key" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_enqueued_team_limit_applies_to_batch_submission():
rate_limiter, local_cache = _enqueued_rate_limiter()
afile_content_mock, patches = _enqueued_batch_patches()
team_user = UserAPIKeyAuth(
api_key="sk-enqueued-team-key",
models=["*"],
team_id="team-enqueued-batch",
team_metadata={"batch_enqueued_token_limit": 10},
)
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
with pytest.raises(HTTPException) as exc:
await rate_limiter.async_pre_call_hook(
user_api_key_dict=team_user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
assert exc.value.status_code == 429
assert "Batch enqueued token limit exceeded for team: team-enqueued-batch" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_disable_flag_still_skips_batch_processing_with_enqueued_limits():
rate_limiter, local_cache = _enqueued_rate_limiter()
afile_content_mock, _ = _enqueued_batch_patches()
opted_in_user = UserAPIKeyAuth(
api_key="sk-enqueued-disabled",
models=["*"],
metadata={"batch_enqueued_token_limit": 10},
)
with (
patch("litellm.proxy.proxy_server.general_settings", {"disable_batch_input_file_rate_limiting": True}),
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch("litellm.afile_content", new=afile_content_mock),
):
data = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
result = await rate_limiter.async_pre_call_hook(
user_api_key_dict=opted_in_user,
cache=local_cache,
data=data,
call_type="acreate_batch",
)
assert result is data
afile_content_mock.assert_not_awaited()

147
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -5858,3 +5858,150 @@ async def test_conflicting_token_limits_cannot_bypass_tpm_reservation():
)
assert exc_info.value.status_code == 429
# ---------------------------------------------------------------------------
# LIT-5273: batch enqueued-token reservations in the post-call hooks
# ---------------------------------------------------------------------------
def _enqueued_test_handler() -> _PROXY_MaxParallelRequestsHandler:
return _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache(default_in_memory_ttl=60)))
def _batch_response(batch_id: str, status: str):
from types import SimpleNamespace
return SimpleNamespace(id=batch_id, status=status, object="batch")
@pytest.mark.asyncio
async def test_success_hook_persists_batch_enqueued_reservation_and_refunds_on_completion():
from litellm.proxy.hooks.batch_enqueued_tokens import (
BatchEnqueuedTokenOverLimit,
BatchEnqueuedTokenReservation,
BatchEnqueuedTokenScope,
)
handler = _enqueued_test_handler()
store = handler.batch_enqueued_token_store
scope = BatchEnqueuedTokenScope(key="api_key", value="hashed-enqueued-key", limit=100)
user = UserAPIKeyAuth(api_key="hashed-enqueued-key")
reservation = await store.reserve(tokens=60, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
get_or_create_request_stash().batch_enqueued_reservation = reservation
await handler.async_post_call_success_hook(
data={}, user_api_key_dict=user, response=_batch_response("batch_enq_1", "validating")
)
assert get_request_stash().batch_enqueued_reservation is None
assert isinstance(await store.reserve(tokens=50, scopes=(scope,)), BatchEnqueuedTokenOverLimit)
await handler.async_post_call_success_hook(
data={}, user_api_key_dict=user, response=_batch_response("batch_enq_1", "completed")
)
refill = await store.reserve(tokens=40, scopes=(scope,))
assert isinstance(refill, BatchEnqueuedTokenReservation)
await handler.async_post_call_success_hook(
data={}, user_api_key_dict=user, response=_batch_response("batch_enq_1", "completed")
)
assert isinstance(await store.reserve(tokens=70, scopes=(scope,)), BatchEnqueuedTokenOverLimit)
@pytest.mark.asyncio
async def test_success_hook_refunds_batch_enqueued_reservation_on_cancellation():
from litellm.proxy.hooks.batch_enqueued_tokens import (
BatchEnqueuedTokenReservation,
BatchEnqueuedTokenScope,
)
handler = _enqueued_test_handler()
store = handler.batch_enqueued_token_store
scope = BatchEnqueuedTokenScope(key="team", value="team-enqueued", limit=100)
user = UserAPIKeyAuth(api_key="hashed-enqueued-key", team_id="team-enqueued")
reservation = await store.reserve(tokens=90, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
get_or_create_request_stash().batch_enqueued_reservation = reservation
await handler.async_post_call_success_hook(
data={}, user_api_key_dict=user, response=_batch_response("batch_enq_2", "validating")
)
await handler.async_post_call_success_hook(
data={}, user_api_key_dict=user, response=_batch_response("batch_enq_2", "cancelling")
)
assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenReservation)
@pytest.mark.asyncio
async def test_success_hook_refunds_on_provider_cased_terminal_status():
from litellm.proxy.hooks.batch_enqueued_tokens import (
BatchEnqueuedTokenOverLimit,
BatchEnqueuedTokenReservation,
BatchEnqueuedTokenScope,
)
handler = _enqueued_test_handler()
store = handler.batch_enqueued_token_store
scope = BatchEnqueuedTokenScope(key="api_key", value="hashed-enqueued-key", limit=100)
user = UserAPIKeyAuth(api_key="hashed-enqueued-key")
reservation = await store.reserve(tokens=90, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
get_or_create_request_stash().batch_enqueued_reservation = reservation
await handler.async_post_call_success_hook(
data={}, user_api_key_dict=user, response=_batch_response("batch_enq_cased", "InProgress")
)
assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenOverLimit)
await handler.async_post_call_success_hook(
data={}, user_api_key_dict=user, response=_batch_response("batch_enq_cased", "Completed")
)
assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenReservation)
@pytest.mark.asyncio
async def test_failure_hook_refunds_stashed_batch_enqueued_reservation():
from litellm.proxy.hooks.batch_enqueued_tokens import (
BatchEnqueuedTokenReservation,
BatchEnqueuedTokenScope,
)
handler = _enqueued_test_handler()
store = handler.batch_enqueued_token_store
scope = BatchEnqueuedTokenScope(key="api_key", value="hashed-failing-key", limit=100)
user = UserAPIKeyAuth(api_key="hashed-failing-key")
reservation = await store.reserve(tokens=80, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
get_or_create_request_stash().batch_enqueued_reservation = reservation
await handler.async_post_call_failure_hook(
request_data={}, original_exception=Exception("guardrail rejected"), user_api_key_dict=user
)
assert get_request_stash().batch_enqueued_reservation is None
assert isinstance(await store.reserve(tokens=100, scopes=(scope,)), BatchEnqueuedTokenReservation)
@pytest.mark.asyncio
async def test_success_hook_leaves_stash_untouched_for_non_batch_responses():
from litellm.proxy.hooks.batch_enqueued_tokens import (
BatchEnqueuedTokenReservation,
BatchEnqueuedTokenScope,
)
handler = _enqueued_test_handler()
store = handler.batch_enqueued_token_store
scope = BatchEnqueuedTokenScope(key="api_key", value="hashed-chat-key", limit=100)
user = UserAPIKeyAuth(api_key="hashed-chat-key")
reservation = await store.reserve(tokens=10, scopes=(scope,))
assert isinstance(reservation, BatchEnqueuedTokenReservation)
get_or_create_request_stash().batch_enqueued_reservation = reservation
await handler.async_post_call_success_hook(
data={}, user_api_key_dict=user, response=ModelResponse(usage=Usage(total_tokens=5))
)
assert get_request_stash().batch_enqueued_reservation == reservation

201
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -15519,6 +15519,207 @@ async def test_regenerate_key_output_token_estimate_lowered_rejected_for_non_adm
assert "Only proxy admins can set" in str(exc.value.detail)
_BATCH_LIMIT = "batch_enqueued_token_limit"
@pytest.mark.parametrize(
"label, request_body, existing_metadata, allowed",
[
("set on a key with none stored", {"metadata": {_BATCH_LIMIT: 50000}}, None, False),
("raised above the stored limit", {"metadata": {_BATCH_LIMIT: 200000}}, {_BATCH_LIMIT: 100000}, False),
("cleared by replacing the blob", {"metadata": {}}, {_BATCH_LIMIT: 100000}, False),
("resent unchanged", {"metadata": {_BATCH_LIMIT: 100000}}, {_BATCH_LIMIT: 100000}, True),
("left untouched", {}, {_BATCH_LIMIT: 100000}, True),
],
)
def test_batch_enqueued_token_limit_admin_gate_matrix(label, request_body, existing_metadata, allowed):
"""A non-admin may only leave a key's stored batch enqueued-token limit as it is.
When set, the limit replaces the standard RPM/TPM checks for batch
submissions, so a key holder writing it would pick their own batch quota.
Resending the stored value is what the edit form produces on every save
and has to stay allowed.
"""
from litellm.proxy.auth.auth_utils import (
enforce_batch_enqueued_token_limit_is_admin_only,
)
def _call(caller):
enforce_batch_enqueued_token_limit_is_admin_only(
data=UpdateKeyRequest(key="sk-1", **request_body),
existing_metadata=existing_metadata,
user_api_key_dict=caller,
entity="key",
)
non_admin = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-non-admin",
user_id="alice",
)
if allowed:
_call(non_admin)
else:
with pytest.raises(HTTPException) as exc:
_call(non_admin)
assert exc.value.status_code == 403
assert "Only proxy admins can set" in str(exc.value.detail)
_call(
UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin",
)
)
@pytest.mark.asyncio
async def test_generate_key_batch_enqueued_token_limit_rejected_for_non_admin():
"""A non-admin self-minting a key with the limit would replace the standard
batch RPM/TPM checks with a cap of their own choosing."""
with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()):
with pytest.raises(HTTPException) as exc:
await _common_key_generation_helper(
data=GenerateKeyRequest(metadata={_BATCH_LIMIT: 100000}, rpm_limit=2),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-alice",
user_id="alice",
),
litellm_changed_by=None,
team_table=None,
)
assert int(getattr(exc.value, "status_code", 0)) == 403
assert "Only proxy admins can set" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_update_key_batch_enqueued_token_limit_raised_rejected_for_non_admin(monkeypatch):
"""/key/update is reachable by the key's own holder, so the gate has to
fire inside the update path itself rather than only at generation."""
from litellm.proxy.management_endpoints.key_management_endpoints import (
update_key_fn,
)
token = "d1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
_wire_update_key_fn(monkeypatch, _estimate_key_row(token, {_BATCH_LIMIT: 100000}))
mock_request = MagicMock()
mock_request.query_params = {}
with pytest.raises(ProxyException) as exc:
await update_key_fn(
request=mock_request,
data=UpdateKeyRequest(key=token, metadata={_BATCH_LIMIT: 10**12}),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-internal",
user_id="internal_user",
),
litellm_changed_by=None,
)
assert str(exc.value.code) == "403"
assert "Only proxy admins can set" in str(exc.value.message)
@pytest.mark.asyncio
async def test_update_key_batch_enqueued_token_limit_unchanged_allows_non_admin_edit(monkeypatch):
"""The edit form resends every field it renders, so gating on presence
would 403 a key owner renaming a key that carries an admin-set limit."""
from litellm.proxy.management_endpoints.key_management_endpoints import (
update_key_fn,
)
token = "e1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
_wire_update_key_fn(monkeypatch, _estimate_key_row(token, {_BATCH_LIMIT: 100000}))
mock_request = MagicMock()
mock_request.query_params = {}
result = await update_key_fn(
request=mock_request,
data=UpdateKeyRequest(key=token, key_alias="my-alias", metadata={_BATCH_LIMIT: 100000}),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-internal",
user_id="internal_user",
),
litellm_changed_by=None,
)
assert result is not None
@pytest.mark.asyncio
async def test_regenerate_key_batch_enqueued_token_limit_rejected_for_non_admin():
"""/key/regenerate runs the request body through prepare_key_update_data
exactly as an update does, so it is a third write path into the field."""
from litellm.proxy._types import RegenerateKeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import (
_execute_virtual_key_regeneration,
)
token = "f1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
key_in_db = LiteLLM_VerificationToken(
token=token,
user_id="internal_user",
metadata={_BATCH_LIMIT: 100000},
)
with pytest.raises(HTTPException) as exc:
await _execute_virtual_key_regeneration(
prisma_client=AsyncMock(),
key_in_db=key_in_db,
hashed_api_key=token,
key="sk-original",
data=RegenerateKeyRequest(key="sk-original", metadata={_BATCH_LIMIT: 10**12}),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-internal",
user_id="internal_user",
),
litellm_changed_by=None,
user_api_key_cache=MagicMock(),
proxy_logging_obj=MagicMock(),
)
assert exc.value.status_code == 403
assert "Only proxy admins can set" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_bulk_key_update_batch_enqueued_token_limit_rejected_for_non_admin():
"""Bulk team-key updates run through _process_single_key_update, not
/key/update's validator, so the gate must also live on that path."""
from litellm.proxy.management_endpoints.key_management_endpoints import (
_process_single_key_update,
)
token = "a2b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
existing = _estimate_key_row(token, {_BATCH_LIMIT: 100000})
with pytest.raises(HTTPException) as exc:
await _process_single_key_update(
update_key_request=UpdateKeyRequest(key=token, metadata={_BATCH_LIMIT: 10**12}),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-internal",
user_id="internal_user",
),
litellm_changed_by=None,
prisma_client=None,
user_api_key_cache=MagicMock(),
proxy_logging_obj=MagicMock(),
llm_router=None,
existing_key_row=existing,
)
assert exc.value.status_code == 403
assert "Only proxy admins can set" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_stamps_settings_updated_at():
"""Regenerate rewrites the key's config, so it must move settings_updated_at."""

57
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -11709,6 +11709,63 @@ async def test_new_team_output_token_estimate_rejected_for_non_admin():
assert "on a team" in str(exc.value.message)
_TEAM_BATCH_LIMIT = "batch_enqueued_token_limit"
@pytest.mark.asyncio
async def test_update_team_batch_enqueued_token_limit_raised_rejected_for_team_admin():
"""_verify_team_access admits a team admin, so the gate has to fire inside
update_team itself to keep the team's batch quota admin-owned."""
import contextlib
from unittest.mock import Mock
from fastapi import Request
from litellm.proxy._types import UpdateTeamRequest
from litellm.proxy.management_endpoints.team_endpoints import update_team
with contextlib.ExitStack() as stack:
_wire_update_team(stack, {_TEAM_BATCH_LIMIT: 100000})
with pytest.raises(ProxyException) as exc:
await update_team(
data=UpdateTeamRequest(team_id="test_team_id", metadata={_TEAM_BATCH_LIMIT: 10**12}),
http_request=Mock(spec=Request),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-team-admin",
user_id="team-admin",
),
)
assert str(exc.value.code) == "403"
assert "on a team" in str(exc.value.message)
@pytest.mark.asyncio
async def test_new_team_batch_enqueued_token_limit_rejected_for_non_admin():
"""/team/new is the other write path into the same stored metadata."""
from unittest.mock import Mock
from fastapi import Request
from litellm.proxy._types import NewTeamRequest
from litellm.proxy.management_endpoints.team_endpoints import new_team
with pytest.raises(ProxyException) as exc:
await new_team(
data=NewTeamRequest(team_alias="t", metadata={_TEAM_BATCH_LIMIT: 100000}),
http_request=Mock(spec=Request),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-alice",
user_id="alice",
),
)
assert str(exc.value.code) == "403"
assert "on a team" in str(exc.value.message)
@pytest.mark.asyncio
async def test_get_team_daily_activity_aggregated_scopes_and_flags(mock_db_client):
"""The aggregated endpoint must apply the same non-admin key scoping as the