diff --git a/litellm/constants.py b/litellm/constants.py index 03955faf59d..fc3dcddab02 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2349,10 +2349,12 @@ PTU_PRUNE_SKEW_GRACE_SECONDS: Final[int] = 300 # 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. +# Key/team metadata field that opts batches into enqueued-token limiting, and its +# per-model project counterpart. Only proxy admins may write them. A key or team limit +# replaces the standard batch rate-limit checks. A project limit on the model the batch +# routes to replaces the per-minute checks, and daily caps still apply. BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit" +MODEL_BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "model_batch_enqueued_token_limit" MAX_BATCH_FILE_RECORDS_KEY: Final = "max_batch_file_records" MAX_BATCH_FILE_UPLOADS_PER_DAY_KEY: Final = "max_batch_file_uploads_per_day" MAX_FILE_DOWNLOADS_PER_MINUTE_KEY: Final = "max_file_downloads_per_minute" diff --git a/litellm/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/litellm/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 0278e8c8c0c..b1f87f962bc 100644 --- a/litellm/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/litellm/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -23,10 +23,13 @@ from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.constants import MODEL_BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import delete_cached_project_object +from litellm.proxy.auth.auth_utils import BatchLimitRequest, enforce_batch_limits_are_admin_only from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.guardrails.guardrail_helpers import validate_guardrail_names +from litellm.proxy.hooks.batch_enqueued_tokens import model_batch_enqueued_token_limits_are_valid from litellm.proxy.management.teams.authz import TEAM_OR_ORG_ADMIN, TeamAccess, is_team_admin, is_team_member from litellm.proxy.management.teams.dependencies import get_team_access from litellm.proxy.management.users.service import org_admin_org_ids @@ -118,6 +121,32 @@ def _jsonified(prisma_client: PrismaClient, payload: dict[str, object]) -> dict[ return jsonified +def _check_project_batch_limits( + data: BatchLimitRequest, + existing_metadata: Mapping[str, object] | None, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + enforce_batch_limits_are_admin_only( + data=data, + existing_metadata=existing_metadata, + user_api_key_dict=user_api_key_dict, + entity="project", + ) + if not model_batch_enqueued_token_limits_are_valid(data.metadata): + raise HTTPException( + status_code=400, + detail={ + "error": f"{MODEL_BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY} must map each model name to a positive integer" + }, + ) + + +def _stored_metadata_with( + request: BatchLimitRequest, stored_metadata: Mapping[str, object] +) -> dict[str, object] | None: + return None if request.metadata is None else {**stored_metadata, **request.metadata} + + async def _check_user_permission_for_project( user_api_key_dict: UserAPIKeyAuth, team_id: str | None, @@ -606,6 +635,11 @@ async def new_project( "error": f"Only admins or team admins can create projects. Your role is {user_api_key_dict.user_role}" }, ) + _check_project_batch_limits( + data=data, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + ) await validate_guardrail_names( None, data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict @@ -845,6 +879,18 @@ async def update_project( status_code=403, detail={"error": "Cannot reassign project to a team you are not an admin of"}, ) + stored_metadata: Final[dict[str, object]] = ( + existing_project.metadata if isinstance(existing_project.metadata, dict) else {} + ) + merged_metadata: Final = None if data.sent_metadata else _stored_metadata_with(data, stored_metadata) + merged_request: Final = ( + data if merged_metadata is None else data.model_copy(update={"metadata": merged_metadata}) + ) + _check_project_batch_limits( + data=merged_request, + existing_metadata=stored_metadata, + user_api_key_dict=user_api_key_dict, + ) await validate_guardrail_names( None, data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict @@ -863,7 +909,7 @@ async def update_project( quota_policy: Final = _project_model_quota_policy(general_settings_view()) if quota_policy is not None: _raise_on_missing_project_model_quota_on_update( - data, + merged_request, existing_project, _router_access_group_names(llm_router), quota_policy, @@ -872,7 +918,7 @@ async def update_project( # Prepare update data await validate_scalar_budget_attachment(prisma_client, data.budget_id) - update_data = _jsonified(prisma_client, data.model_dump(exclude_none=True, exclude={"project_id"})) + update_data = _jsonified(prisma_client, merged_request.model_dump(exclude_none=True, exclude={"project_id"})) update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name # Handle budget updates diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 53601c1670f..552607e614c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -15,8 +15,10 @@ from pydantic import ( Field, Json, JsonValue, + ModelWrapValidatorHandler, NonNegativeInt, PositiveInt, + PrivateAttr, TypeAdapter, ValidationError, ValidatorFunctionWrapHandler, @@ -28,6 +30,7 @@ from typing_extensions import ( NotRequired, ReadOnly, Required, + Self, TypedDict, ) @@ -3518,6 +3521,7 @@ class UpdateProjectRequest(LiteLLM_BudgetTable): blocked: bool | None = None budget_id: str | None = None object_permission: LiteLLM_ObjectPermissionBase | None = None + _sent_metadata: bool = PrivateAttr(default=True) @model_validator(mode="before") @classmethod @@ -3533,6 +3537,18 @@ class UpdateProjectRequest(LiteLLM_BudgetTable): values.pop(field) return values + @model_validator(mode="wrap") + @classmethod + def remember_whether_metadata_was_sent(cls, values: object, handler: ModelWrapValidatorHandler[Self]) -> Self: + sent_metadata: Final = not isinstance(values, Mapping) or "metadata" in values + request: Final = handler(values) + request._sent_metadata = sent_metadata + return request + + @property + def sent_metadata(self) -> bool: + return self._sent_metadata + class DeleteProjectRequest(LiteLLMPydanticObjectBase): """Request model for DELETE /project/delete""" diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 827dbeee9b7..259c46a9dc1 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -21,6 +21,7 @@ from litellm.constants import ( EMPTY_MAPPING, INVALID_VIRTUAL_KEY_ERROR_MARKER, MINIMUM_CUSTOM_KEY_LENGTH, + MODEL_BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, STANDARD_CUSTOMER_ID_HEADERS, ) from litellm.litellm_core_utils.clientside_credential_handler import ( @@ -1928,15 +1929,16 @@ def enforce_batch_limits_are_admin_only( data: BatchLimitRequest, existing_metadata: Mapping[str, object] | None, user_api_key_dict: UserAPIKeyAuth, - entity: Literal["key", "team"], + entity: Literal["key", "team", "project"], ) -> None: - """Only a proxy admin may change a key or team's batch limits. + """Only a proxy admin may change a key, team, or project's batch limits. Every key in ``ADMIN_ONLY_BATCH_LIMIT_METADATA_KEYS`` caps what the holder may do with batches, so a holder-writable copy would let a caller lift their own quota, and the S3 bucket list grants access to buckets with the proxy's - own credentials. Gated on the resulting value rather than on presence, so a form - resending the stored value stays a no-op. + own credentials. A project only enforces ``model_batch_enqueued_token_limit``, + so that is the one key gated there. 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 @@ -1944,8 +1946,11 @@ def enforce_batch_limits_are_admin_only( requested: Final[Mapping[str, object]] = ( (data.metadata or EMPTY_MAPPING) if "metadata" in data.model_fields_set else stored ) + admin_only_keys: Final[tuple[str, ...]] = ( + (MODEL_BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY,) if entity == "project" else _ADMIN_ONLY_METADATA_KEYS + ) changed: Final = next( - (key for key in _ADMIN_ONLY_METADATA_KEYS if requested.get(key) != stored.get(key)), + (key for key in admin_only_keys if requested.get(key) != stored.get(key)), None, ) if changed is None: diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index e29855056fa..d408c4d9c78 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -1,7 +1,8 @@ """ Enqueued-token accounting for batch submissions. -Opt-in via admin-set ``batch_enqueued_token_limit`` in key or team metadata: batch +Opt-in via admin-set ``batch_enqueued_token_limit`` in key or team metadata, or a +per-model ``model_batch_enqueued_token_limit`` in project 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 @@ -9,20 +10,28 @@ the reservation is refunded when the batch reaches a terminal state """ import asyncio +import itertools import logging import math +import re import time import uuid -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass, field +from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias -from pydantic import ConfigDict, Field, TypeAdapter, ValidationError +from pydantic import ConfigDict, Field, StrictInt, TypeAdapter, ValidationError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import log_redis_failure -from litellm.constants import BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, BATCH_ENQUEUED_TOKEN_TTL_SECONDS +from litellm.constants import ( + BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, + BATCH_ENQUEUED_TOKEN_TTL_SECONDS, + MODEL_BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, +) +from litellm.litellm_core_utils.model_latest_aliases import expand_model_latest_alias from litellm.proxy._types import UserAPIKeyAuth from litellm.redis.stores.rate_limits import RATE_LIMITS_TARGET from litellm.types.llms.base import LiteLLMBaseModel @@ -39,7 +48,7 @@ BATCH_ENQUEUED_REFUND_STATUSES: Final[frozenset[str]] = frozenset( {"completed", "complete", "failed", "expired", "cancelled", "cancelling"} ) -ScopeKey: TypeAlias = Literal["api_key", "team"] +ScopeKey: TypeAlias = Literal["api_key", "team", "project_model"] RESERVE_ENQUEUED_TOKENS_SCRIPT: Final = """ local amount = tonumber(ARGV[1]) @@ -81,6 +90,10 @@ class BatchEnqueuedTokenScope: key: ScopeKey value: str limit: int + tokens: int | None = None + + def charge(self, batch_tokens: int) -> int: + return batch_tokens if self.tokens is None else self.tokens ReservationBackend: TypeAlias = Literal["redis", "memory"] @@ -103,11 +116,42 @@ class BatchEnqueuedTokenOverLimit: BatchEnqueuedTokenOutcome: TypeAlias = BatchEnqueuedTokenReservation | BatchEnqueuedTokenOverLimit + +@dataclass(frozen=True, slots=True) +class _StoredReservation: + """A reservation record as stored, with project scopes apart so pods that predate them still refund the rest.""" + + tokens: int + scopes: tuple[BatchEnqueuedTokenScope, ...] + backend: ReservationBackend = "redis" + owner: str = "" + project_scopes: tuple[BatchEnqueuedTokenScope, ...] = () + + @classmethod + def of(cls, reservation: BatchEnqueuedTokenReservation) -> "_StoredReservation": + return cls( + tokens=reservation.tokens, + scopes=tuple(scope for scope in reservation.scopes if scope.key != "project_model"), + backend=reservation.backend, + owner=reservation.owner, + project_scopes=tuple(scope for scope in reservation.scopes if scope.key == "project_model"), + ) + + def restored(self) -> BatchEnqueuedTokenReservation: + return BatchEnqueuedTokenReservation( + tokens=self.tokens, scopes=(*self.scopes, *self.project_scopes), backend=self.backend, owner=self.owner + ) + + _LIMIT_ADAPTER: Final = TypeAdapter(Annotated[int, Field(gt=0)]) +_MODEL_LIMITS_ADAPTER: Final[TypeAdapter[Mapping[str, int]]] = TypeAdapter( + Mapping[str, Annotated[StrictInt, Field(gt=0)]] +) +_NO_MODEL_LIMITS: Final[Mapping[str, int]] = MappingProxyType({}) _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) +_RESERVATION_ADAPTER: Final = TypeAdapter(_StoredReservation) _LOCAL_RECORD_ADAPTER: Final = TypeAdapter(object) @@ -148,6 +192,112 @@ def resolve_batch_enqueued_token_scopes( return tuple(scope for scope in candidates if scope is not None) +def _read_project_model_limits(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, int]: + metadata: Final[Mapping[str, object]] = user_api_key_dict.project_metadata or _NO_MODEL_LIMITS + raw: Final = metadata.get(MODEL_BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY) + if raw is None or not user_api_key_dict.project_id: + return _NO_MODEL_LIMITS + try: + return MappingProxyType(_MODEL_LIMITS_ADAPTER.validate_python(raw)) + except ValidationError: + verbose_proxy_logger.warning( + "Ignoring invalid %s value %r; expected a mapping of model name to positive integer", + MODEL_BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, + raw, + ) + return _NO_MODEL_LIMITS + + +def model_batch_enqueued_token_limits_are_valid(metadata: Mapping[str, object] | None) -> bool: + raw: Final = (metadata or _NO_MODEL_LIMITS).get(MODEL_BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY) + if raw is None: + return True + try: + _MODEL_LIMITS_ADAPTER.validate_python(raw) + except ValidationError: + return False + return True + + +def project_has_model_batch_enqueued_token_limits(user_api_key_dict: UserAPIKeyAuth) -> bool: + return bool(_read_project_model_limits(user_api_key_dict)) + + +_DATED_SNAPSHOT_SUFFIX: Final = re.compile(r"-\d{4}-\d{2}-\d{2}$") + + +def _tightest_limit_by_group(limits: Mapping[str, int], model_group_of: Callable[[str], str]) -> Mapping[str, int]: + resolved: Final = tuple((model_group_of(name), limit) for name, limit in limits.items()) + return MappingProxyType({group: min(limit for g, limit in resolved if g == group) for group, _ in resolved}) + + +def _model_groups(model: str, model_group_of: Callable[[str], str]) -> frozenset[str]: + names: Final = (*expand_model_latest_alias(model), _DATED_SNAPSHOT_SUFFIX.sub("", model)) + return frozenset(model_group_of(candidate) for candidate in names) + + +def _limited_model_charges( + limits: Mapping[str, int], + tokens_by_model: Mapping[str, int], + model_group_of: Callable[[str], str], +) -> Iterator[tuple[str, int, int]]: + group_limits: Final = _tightest_limit_by_group(limits, model_group_of) + return itertools.chain.from_iterable( + ( + (group, group_limits[group], tokens) + for group in _model_groups(model, model_group_of) + if group in group_limits + ) + for model, tokens in tokens_by_model.items() + ) + + +def project_limits_every_model( + user_api_key_dict: UserAPIKeyAuth, + models: Iterable[str], + model_group_of: Callable[[str], str], +) -> bool: + group_limits: Final = _tightest_limit_by_group(_read_project_model_limits(user_api_key_dict), model_group_of) + return all(not group_limits.keys().isdisjoint(_model_groups(model, model_group_of)) for model in models) + + +def resolve_project_model_batch_enqueued_token_scopes( + user_api_key_dict: UserAPIKeyAuth, + row_tokens_by_model: Mapping[str, int], + routing_model: str | None, + batch_tokens: int, + model_group_of: Callable[[str], str], +) -> tuple[BatchEnqueuedTokenScope, ...]: + """One scope per limited model group a batch touches, under the tightest limit keyed by any of its names. + + The routing model is charged the whole batch, because rows run on it whatever model they name, and + every other model a row names is charged that row's tokens, so neither can hide behind the other. + """ + limits: Final = _read_project_model_limits(user_api_key_dict) + row_charges: Final = tuple(_limited_model_charges(limits, row_tokens_by_model, model_group_of)) + routing_charges: Final = tuple( + _limited_model_charges( + limits, + MappingProxyType({routing_model: batch_tokens}) if routing_model else _NO_MODEL_LIMITS, + model_group_of, + ) + ) + charges: Final = (*row_charges, *routing_charges) + groups: Final = sorted(frozenset(group for group, _, _ in charges)) + return tuple( + BatchEnqueuedTokenScope( + key="project_model", + value=f"{user_api_key_dict.project_id}:{group}", + limit=min(limit for charge_group, limit, _ in charges if charge_group == group), + tokens=max( + sum(tokens for charge_group, _, tokens in row_charges if charge_group == group), + max((tokens for charge_group, _, tokens in routing_charges if charge_group == group), default=0), + ), + ) + for group in groups + ) + + def canonical_provider_batch_id(batch_id: str) -> str: from litellm.proxy.openai_files_endpoints.common_utils import ( get_batch_id_from_unified_batch_id, @@ -282,7 +432,7 @@ class BatchEnqueuedTokenStore: try: raw_result: Final = await reserve_script( (self._counter_key(scope),), - (tokens, BATCH_ENQUEUED_TOKEN_TTL_SECONDS, scope.limit), + (scope.charge(tokens), BATCH_ENQUEUED_TOKEN_TTL_SECONDS, scope.limit), ) return _RESERVE_RESULT_ADAPTER.validate_python(raw_result) except Exception: @@ -310,7 +460,7 @@ class BatchEnqueuedTokenStore: scopes: tuple[BatchEnqueuedTokenScope, ...], ) -> None: for scope in scopes: - await refund_script((self._counter_key(scope),), (tokens,)) + await refund_script((self._counter_key(scope),), (scope.charge(tokens),)) async def _reserve_in_memory( self, @@ -322,10 +472,10 @@ class BatchEnqueuedTokenStore: 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: + if current + scope.charge(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) + await self._set_local_counter(scope, current + scope.charge(tokens), span) return BatchEnqueuedTokenReservation( tokens=tokens, scopes=scopes, backend="memory", owner=self._owner_token, reserved_at_monotonic=started ) @@ -349,7 +499,7 @@ class BatchEnqueuedTokenStore: async with self._lock: for scope in reservation.scopes: current = await self._get_local_counter(scope, litellm_parent_otel_span) - remaining = current - reservation.tokens + remaining = current - scope.charge(reservation.tokens) if remaining <= 0: self.internal_usage_cache.dual_cache.in_memory_cache.delete_cache(key=self._counter_key(scope)) else: @@ -376,7 +526,7 @@ class BatchEnqueuedTokenStore: reservation: BatchEnqueuedTokenReservation, litellm_parent_otel_span: "Span | None" = None, ) -> None: - serialized: Final = _RESERVATION_ADAPTER.dump_json(reservation).decode("utf-8") + serialized: Final = _RESERVATION_ADAPTER.dump_json(_StoredReservation.of(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: @@ -422,8 +572,8 @@ class BatchEnqueuedTokenStore: return None try: if isinstance(raw, (str, bytes)): - return _RESERVATION_ADAPTER.validate_json(raw) - return _RESERVATION_ADAPTER.validate_python(raw) + return _RESERVATION_ADAPTER.validate_json(raw).restored() + return _RESERVATION_ADAPTER.validate_python(raw).restored() except ValidationError: verbose_proxy_logger.warning("Discarding malformed enqueued-token reservation record for %s", batch_id) return None diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 625e1e2df7e..d8277f2ee38 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -18,7 +18,7 @@ Quick summary: """ import json -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, TypeAlias @@ -56,7 +56,10 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( BatchEnqueuedTokenOverLimit, BatchEnqueuedTokenReservation, BatchEnqueuedTokenScope, + project_has_model_batch_enqueued_token_limits, + project_limits_every_model, resolve_batch_enqueued_token_scopes, + resolve_project_model_batch_enqueued_token_scopes, ) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROJECT_ITPM_DESCRIPTOR_KEY, @@ -151,6 +154,19 @@ class BatchFileUsage(LiteLLMBaseModel): per_model_usage: dict[str, dict[str, int]] = Field(default_factory=dict) +def _rows_were_rewritten_on_upload(input_file_id: str) -> bool: + from litellm.proxy.openai_files_endpoints.common_utils import is_base64_encoded_unified_file_id + + return bool(is_base64_encoded_unified_file_id(input_file_id)) + + +def _deployment_model_name_via_proxy_router(model: str) -> str | None: + from litellm.proxy.proxy_server import llm_router + + deployment: Final = None if llm_router is None else llm_router.registry().deployment(model) + return None if deployment is None else deployment.model_name + + class _PROXY_BatchRateLimiter(CustomLogger): """ Rate limiter for batch API requests. @@ -165,6 +181,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): internal_usage_cache: InternalUsageCache, parallel_request_limiter: ParallelRequestLimiter, time_provider: Callable[[], datetime] | None = None, + deployment_model_name: Callable[[str], str | None] = _deployment_model_name_via_proxy_router, ): """ Initialize the batch rate limiter. @@ -176,10 +193,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): internal_usage_cache: Cache for storing rate limit data (auto-injected) parallel_request_limiter: Existing rate limiter to integrate with (needs custom injection) time_provider: Clock used for rate limit reset times (defaults to ``datetime.now``) + deployment_model_name: Model name of the deployment a batch names by id, else None """ self.internal_usage_cache = internal_usage_cache self.parallel_request_limiter = parallel_request_limiter self._time_provider: Final = time_provider or datetime.now + self._deployment_model_name: Final = deployment_model_name self._warned_unsupported_model_skip = False def _get_file_bound_batch_model(self, data: dict) -> str | None: @@ -226,11 +245,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): if file_bound_model: return file_bound_model - model: Final = data.get("model") - if isinstance(model, str) and model: - return model - - return None + return _batch_create_routing_model(_BATCH_BODY_ADAPTER.validate_python(data)) def _resolve_batch_provider(self, batch_model: str | None) -> str | None: """Resolve the provider from the deployment that serves ``batch_model``. @@ -622,8 +637,10 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) -> 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 + Runs instead of the per-minute counter charge when the key or team opted + in via ``batch_enqueued_token_limit`` metadata, or the project limits the + batch's routing model via ``model_batch_enqueued_token_limit``; otherwise it + runs alongside the normal counters. 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. @@ -639,6 +656,27 @@ class _PROXY_BatchRateLimiter(CustomLogger): case BatchEnqueuedTokenReservation(): get_or_create_request_stash().batch_enqueued_reservation = outcome + async def _refund_enqueued_reservation_if_rejected( + self, user_api_key_dict: UserAPIKeyAuth, check: Awaitable[None] + ) -> None: + """Run a check that follows the enqueued-token reservation, refunding the reservation if it rejects the batch.""" + try: + await check + except HTTPException: + stash: Final = get_or_create_request_stash() + if stash.batch_enqueued_reservation is not None: + await self.parallel_request_limiter.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 + raise + + def _model_group_of(self, model: str) -> str: + named_model: Final = self._deployment_model_name(model) or model + rate_limited: Final = self.parallel_request_limiter.rate_limited_model(named_model) + return rate_limited.group if rate_limited is not None else named_model + def _raise_enqueued_limit_error( self, over_limit: BatchEnqueuedTokenOverLimit, @@ -649,7 +687,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): 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"Batch requires {scope.charge(batch_usage.total_tokens)} tokens " + f"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." ) @@ -773,6 +812,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): data: dict, batch_usage: BatchFileUsage, descriptors: list["RateLimitDescriptor"] | None = None, + include_project_io_limits: bool = True, ) -> None: """ Atomically check + increment rate-limit counters by the batch amounts. @@ -802,9 +842,13 @@ class _PROXY_BatchRateLimiter(CustomLogger): for _d in descriptors ] - project_io_descriptors, project_io_increments = self._create_project_io_descriptors_for_models( - user_api_key_dict=user_api_key_dict, - per_model_usage=batch_usage.per_model_usage, + project_io_descriptors, project_io_increments = ( + self._create_project_io_descriptors_for_models( + user_api_key_dict=user_api_key_dict, + per_model_usage=batch_usage.per_model_usage, + ) + if include_project_io_limits + else ([], []) ) descriptors = [*descriptors, *project_io_descriptors] increments = [*increments, *project_io_increments] @@ -1286,9 +1330,12 @@ 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) + key_and_team_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, has_enqueued_scopes=bool(enqueued_scopes) + data=data, + user_api_key_dict=user_api_key_dict, + has_enqueued_scopes=bool(key_and_team_scopes) + or project_has_model_batch_enqueued_token_limits(user_api_key_dict), ) if should_skip: return data @@ -1314,6 +1361,22 @@ class _PROXY_BatchRateLimiter(CustomLogger): data["_batch_token_count"] = batch_usage.total_tokens data["_batch_request_count"] = batch_usage.request_count + routing_model: Final = self._get_batch_routing_model(data) # pyright: ignore[reportUnknownMemberType] # hook data is the untyped proxy request dict + row_tokens_by_model: Final[Mapping[str, int]] = MappingProxyType( + {model: usage.get("total_tokens", 0) for model, usage in batch_usage.per_model_usage.items()} + if not (isinstance(input_file_id, str) and _rows_were_rewritten_on_upload(input_file_id)) + else {} + ) + enqueued_scopes: Final = ( + *key_and_team_scopes, + *resolve_project_model_batch_enqueued_token_scopes( + user_api_key_dict, + row_tokens_by_model=row_tokens_by_model, + routing_model=routing_model, + batch_tokens=batch_usage.total_tokens, + model_group_of=self._model_group_of, + ), + ) if enqueued_scopes: await self._reserve_batch_enqueued_tokens( user_api_key_dict=user_api_key_dict, @@ -1321,6 +1384,32 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage=batch_usage, scopes=enqueued_scopes, ) + if not key_and_team_scopes: + project_covers_batch: Final = routing_model is not None and project_limits_every_model( + user_api_key_dict, (routing_model, *row_tokens_by_model), self._model_group_of + ) + remaining_descriptors: Final = [ + descriptor + for descriptor in ( + batch_rate_limit_descriptors + if batch_rate_limit_descriptors is not None + else self._create_batch_rate_limit_descriptors( # pyright: ignore[reportUnknownMemberType] # hook data is the untyped proxy request dict + user_api_key_dict=user_api_key_dict, data=data + ) + ) + if not project_covers_batch or descriptor["key"].endswith(BATCH_TPD_DESCRIPTOR_SUFFIX) + ] + if remaining_descriptors or not project_covers_batch: + await self._refund_enqueued_reservation_if_rejected( + user_api_key_dict, + self._check_and_increment_batch_counters( # pyright: ignore[reportUnknownMemberType] # hook data is the untyped proxy request dict + user_api_key_dict=user_api_key_dict, + data=data, + batch_usage=batch_usage, + descriptors=remaining_descriptors, + include_project_io_limits=not project_covers_batch, + ), + ) verbose_proxy_logger.debug("Batch enqueued-token reservation succeeded") return data diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx index 2e39d45e660..5105d6c7220 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx @@ -10,8 +10,9 @@ vi.mock("@/app/(dashboard)/hooks/projects/useUpdateProject", () => ({ useUpdateProject: () => ({ mutate, isPending: false }), })); +const mockUseAuthorized = vi.fn(); vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: () => ({ accessToken: "test-token", userId: "u-1", userRole: "Admin" }), + default: () => mockUseAuthorized(), })); vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ @@ -79,6 +80,7 @@ const variables = () => mutate.mock.calls.at(-1)?.[0] as { projectId: string; pa describe("EditProjectModal submit payload", () => { beforeEach(() => { vi.clearAllMocks(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userId: "u-1", userRole: "Admin" }); }); it("sends exactly the antd payload for an untouched save", async () => { @@ -172,6 +174,102 @@ describe("EditProjectModal submit payload", () => { expect(variables().params.metadata).toStrictEqual({ owner: "platform" }); }); + it("keeps the stored per-model batch enqueued limits when saving after opening Advanced Settings", async () => { + const user = setup(); + renderModal({ + ...project, + metadata: { ...project.metadata, model_batch_enqueued_token_limit: { "gpt-4": 1000 } }, + } as unknown as ProjectResponse); + await screen.findByDisplayValue("My Project"); + + await user.click(screen.getByText("Advanced Settings")); + await screen.findByText("Model-Specific Limits"); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(variables().params.metadata).toStrictEqual({ + owner: "platform", + model_batch_enqueued_token_limit: { "gpt-4": 1000 }, + }); + }); + + it("sends the per-model batch enqueued limit edited in its model limit row", async () => { + const user = setup(); + renderModal({ + ...project, + metadata: { ...project.metadata, model_batch_enqueued_token_limit: { "gpt-4": 1000 } }, + } as unknown as ProjectResponse); + await screen.findByDisplayValue("My Project"); + + await user.click(screen.getByText("Advanced Settings")); + const batchLimit = await screen.findByLabelText("Batch Enqueued Token Limit"); + expect(batchLimit).toHaveValue(1000); + fireEvent.change(batchLimit, { target: { value: "2500" } }); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(variables().params.metadata).toStrictEqual({ + owner: "platform", + model_batch_enqueued_token_limit: { "gpt-4": 2500 }, + }); + }); + + it("drops the per-model batch enqueued limit once its model limit row is removed", async () => { + const user = setup(); + renderModal({ + ...project, + metadata: { ...project.metadata, model_batch_enqueued_token_limit: { "gpt-4": 1000 } }, + } as unknown as ProjectResponse); + await screen.findByDisplayValue("My Project"); + + await user.click(screen.getByText("Advanced Settings")); + await screen.findByText("Model-Specific Limits"); + await user.click(screen.getByRole("button", { name: "Remove model limit 1" })); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(variables().params.metadata).toStrictEqual({ owner: "platform" }); + }); + + it("blocks the save when the per-model batch enqueued limit is not a whole number above 0", async () => { + const user = setup(); + renderModal({ + ...project, + metadata: { ...project.metadata, model_batch_enqueued_token_limit: { "gpt-4": 1000 } }, + } as unknown as ProjectResponse); + await screen.findByDisplayValue("My Project"); + + await user.click(screen.getByText("Advanced Settings")); + const batchLimit = await screen.findByLabelText("Batch Enqueued Token Limit"); + fireEvent.change(batchLimit, { target: { value: "0" } }); + await save(user); + + expect(await screen.findByText("Enter a whole number above 0")).toBeInTheDocument(); + expect(mutate).not.toHaveBeenCalled(); + }); + + it("keeps the stored per-model batch enqueued limit when a team admin removes its model limit row", async () => { + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userId: "u-2", userRole: "Internal User" }); + const user = setup(); + renderModal({ + ...project, + metadata: { ...project.metadata, model_batch_enqueued_token_limit: { "gpt-4": 1000 } }, + } as unknown as ProjectResponse); + await screen.findByDisplayValue("My Project"); + + await user.click(screen.getByText("Advanced Settings")); + expect(await screen.findByLabelText("Batch Enqueued Token Limit")).toBeDisabled(); + await user.click(screen.getByRole("button", { name: "Remove model limit 1" })); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(variables().params).toMatchObject({ model_tpm_limit: {}, model_rpm_limit: {} }); + expect(variables().params.metadata).toStrictEqual({ + owner: "platform", + model_batch_enqueued_token_limit: { "gpt-4": 1000 }, + }); + }); + it("blocks the save when the project name is cleared", async () => { const user = setup(); renderModal(); diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx index 1b4277251b4..33d24d87e06 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx @@ -10,9 +10,11 @@ import { Button } from "@/components/ui/button"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUpdateProject, ProjectUpdateParams } from "@/app/(dashboard)/hooks/projects/useUpdateProject"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { isProxyAdminRole } from "@/utils/roles"; import { ProjectBaseForm } from "./ProjectBaseForm"; import { projectFormSchema, type ProjectFormValues, type ProjectSubmitValues } from "./projectFormSchema"; -import { buildProjectUpdateParams } from "./projectFormUtils"; +import { buildProjectUpdateParams, withStoredBatchEnqueuedTokenLimit } from "./projectFormUtils"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; interface EditProjectModalProps { @@ -27,6 +29,7 @@ const INTERNAL_METADATA_KEYS = new Set([ "model_tpm_limit", "model_itpm_limit", "model_otpm_limit", + "model_batch_enqueued_token_limit", "guardrails", ]); @@ -36,6 +39,7 @@ export const toFormValues = (project: ProjectResponse): ProjectFormValues => { const tpmLimits = (metadataObj.model_tpm_limit ?? {}) as Record; const itpmLimits = (metadataObj.model_itpm_limit ?? {}) as Record; const otpmLimits = (metadataObj.model_otpm_limit ?? {}) as Record; + const batchEnqueuedTokenLimits = (metadataObj.model_batch_enqueued_token_limit ?? {}) as Record; const guardrails = (Array.isArray(metadataObj.guardrails) ? metadataObj.guardrails : []) as string[]; const modelLimits = Array.from( @@ -44,6 +48,7 @@ export const toFormValues = (project: ProjectResponse): ProjectFormValues => { ...Object.keys(tpmLimits), ...Object.keys(itpmLimits), ...Object.keys(otpmLimits), + ...Object.keys(batchEnqueuedTokenLimits), ]), ).map((model) => ({ model, @@ -51,6 +56,7 @@ export const toFormValues = (project: ProjectResponse): ProjectFormValues => { tpm: tpmLimits[model], itpm: itpmLimits[model], otpm: otpmLimits[model], + batchEnqueuedTokens: batchEnqueuedTokenLimits[model], })); const metadata = Object.entries(metadataObj) @@ -74,6 +80,7 @@ function EditProjectForm({ project, onClose, onSuccess }: Omit(null); @@ -58,6 +60,7 @@ export function ProjectBaseForm({ form, advancedOpen, onAdvancedOpenChange }: Pr rpm: undefined, itpm: undefined, otpm: undefined, + batchEnqueuedTokens: undefined, }; const teamIdValue = useWatch({ control: form.control, name: "team_id" }); @@ -292,7 +295,7 @@ export function ProjectBaseForm({ form, advancedOpen, onAdvancedOpenChange }: Pr {modelLimits.fields.map((field, index) => (
)} + + {({ ref, value, onChange, ...control }) => ( + onChange(toOptionalNumber(event.target.value))} + /> + )} +