diff --git a/litellm/constants.py b/litellm/constants.py index 39a49e55f0d..facfc6f7c19 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1766,6 +1766,17 @@ # 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({}) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index a105bf19458..883f986f6fd 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -12,7 +12,12 @@ 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"], diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py new file mode 100644 index 00000000000..32bccca2ab0 --- /dev/null +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -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, + ) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index d6229fb80a6..5b814ad28fd 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -46,9 +46,16 @@ 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 @@ def _should_skip_batch_input_file_processing( 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 @@ def _should_skip_batch_input_file_processing( 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 @@ def _resolve_batch_input_file_fetch_params( 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 @@ async def async_pre_call_hook( 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 @@ async def async_pre_call_hook( 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( diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 2d799ded752..1e65da5b867 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -44,6 +44,13 @@ 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 @@ def __init__( # 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 @@ async def async_post_call_success_hook(self, data: dict, user_api_key_dict: User 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 @@ async def async_post_call_failure_hook( ) 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 diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 5f823c5b20d..71218d6114b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -57,6 +57,7 @@ ) 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) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0211d86b83d..82e22bb5bbf 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -85,7 +85,10 @@ 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") diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 1376bdbed38..53bf9739983 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -17,6 +17,7 @@ import json import os +import re import time from datetime import datetime, timedelta, timezone from typing import Callable @@ -57,7 +58,7 @@ 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 @@ def test_batch_create_over_rpm_returns_mapped_429( ) +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" diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 2dfa7adddea..4b8aa1da002 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -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"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 619f4dcacfe..ac41971a2c8 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -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): diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py new file mode 100644 index 00000000000..edc921a40a3 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py @@ -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) diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index 28624254565..4ef94f0965b 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -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() diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 54366226dfb..fee892ad342 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 0936a2d1439..8e661af8daa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index c20b6062e7a..2d54a391cf0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -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