From 73e17cba7458f1f8a219c05ea2310e93f0ddcb4b Mon Sep 17 00:00:00 2001 From: kirti22jain <31334586+kirti22jain@users.noreply.github.com> Date: Sun, 15 Nov 2026 01:06:53 +0530 Subject: [PATCH] feat(proxy): add per-model project batch enqueued token limits (#51608) * feat(proxy): add per-model project batch enqueued token limits Enqueued batch token limits only read key and team metadata, so a deployment that gives each application its own project under a shared team cannot cap one application's queued batch tokens for a given model. The team limit is shared by every project in it, and JWT or custom-auth callers get a new key scope with every token. Read model_batch_enqueued_token_limit from project metadata, keyed by model like model_tpm_limit. Each limited model group the batch runs is reserved against only its own rows' tokens, rows that name no model count toward the routing model, and the reservation records each scope's charge so the terminal-state refund releases exactly that. Only proxy admins may set it, on project create and update as well as on keys and teams, so a team admin cannot lift their own project's cap. * fix(proxy): charge the routing model and keep daily caps for project batch limits Rows can name any model while the batch runs on the routing deployment, for example LiteLLM-executed batches replace every row's model with the run model, so charging only row models let a batch skip a capped deployment's per-model limit. Charge the routing model the whole batch as well as each other limited row model its own rows' tokens. An enqueued-token reservation replaced every batch counter, so a project limit also skipped the key, team and end-user daily token caps. Keep checking those daily caps after reserving, and refund the reservation if one rejects the batch. * fix(proxy): keep per-minute checks for batches a project allowance does not cover A per-model project allowance replaced every per-minute batch check as soon as any row's model had one. A batch routed to an unlimited model could add a single tiny row naming a limited model and skip the routing model's RPM/TPM and project ITPM/OTPM quotas, although every row runs on the routing model. Waive those checks only when the reservation covers the whole batch, through a key or team allowance or a project allowance on the routing model; otherwise reserve the covered models and run the normal checks too. * test(proxy): cover the project per-model batch cap through custom auth on a live proxy * fix(proxy): keep key and team batch allowances authoritative and harden project limits Key and team enqueued allowances again replace the other batch checks. A project allowance on the routing model keeps only the daily caps. Team admins can only be blocked from the project key, invalid project maps answer 400, the edit modal keeps stored limits, aliases share their tightest limit, file-bound batches charge only the routing model, deployment ids resolve to their model group, and reservation records keep project scopes apart so older pods still refund key and team scopes. * fix(proxy): charge row models for model-bound batch uploads Only a managed-file upload rewrites rows to the deployment model. A model-bound upload keeps each row's model, so its rows still count toward the project caps of the models they name. * fix(proxy): skip per-minute batch limits only when the project caps every model in the batch A project cap on the routing model alone used to switch off the project's per-minute input and output limits for the other models a batch file names. The skip now needs a project cap on every one of those models. -latest names are charged under the cap of the model they point to, the limit map only takes JSON integers, and a /project/update that sends only top-level fields keeps the stored metadata instead of replacing it. * fix(proxy): lift per-minute limits off managed batches under a project cap and edit the cap in the Admin UI * fix(ui): keep the project batch cap read-only for team admins and require a whole number above 0 * fix(ui): narrow project metadata directly when keeping the stored batch cap * fix(proxy): charge dated snapshot rows and header or query routed batches to project caps --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 8 +- .../management_endpoints/project_endpoints.py | 50 ++- litellm/proxy/_types.py | 16 + litellm/proxy/auth/auth_utils.py | 15 +- litellm/proxy/hooks/batch_enqueued_tokens.py | 178 +++++++++- litellm/proxy/hooks/batch_rate_limiter.py | 117 ++++++- .../EditProjectModal.integration.test.tsx | 100 +++++- .../ProjectModals/EditProjectModal.tsx | 19 +- .../ProjectModals/ProjectBaseForm.tsx | 24 +- .../ProjectModals/projectFormSchema.ts | 5 + .../ProjectModals/projectFormUtils.test.ts | 45 ++- .../ProjectModals/projectFormUtils.ts | 30 +- .../src/i18n/locales/en.json | 4 +- ...test_project_batch_enqueued_token_limit.py | 188 +++++++++++ .../test_project_endpoints_prisma.py | 166 ++++++++- .../proxy/hooks/test_batch_enqueued_tokens.py | 241 ++++++++++++- .../proxy/hooks/test_batch_file_validation.py | 317 +++++++++++++++++- .../proxy/hooks/test_batch_rate_limiter.py | 142 +++++++- 18 files changed, 1594 insertions(+), 71 deletions(-) create mode 100644 tests/integration/spend/test_project_batch_enqueued_token_limit.py 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))} + /> + )} +