mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
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>
This commit is contained in:
kirti22jain
GitHub
parent
9fba44d300
commit
73e17cba74
18 changed files with 1594 additions and 71 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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<string, number>;
|
||||
const itpmLimits = (metadataObj.model_itpm_limit ?? {}) as Record<string, number>;
|
||||
const otpmLimits = (metadataObj.model_otpm_limit ?? {}) as Record<string, number>;
|
||||
const batchEnqueuedTokenLimits = (metadataObj.model_batch_enqueued_token_limit ?? {}) as Record<string, number>;
|
||||
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<EditProjectModalP
|
|||
const { t } = useTranslation();
|
||||
const form = useZodForm(projectFormSchema, { defaultValues: toFormValues(project) });
|
||||
const updateMutation = useUpdateProject();
|
||||
const { userRole } = useAuthorized();
|
||||
const [advancedOpen, setAdvancedOpen] = useState(false);
|
||||
const [advancedEverOpened, setAdvancedEverOpened] = useState(false);
|
||||
|
||||
|
|
@ -87,10 +94,12 @@ function EditProjectForm({ project, onClose, onSuccess }: Omit<EditProjectModalP
|
|||
? values
|
||||
: { ...values, guardrails: undefined, modelLimits: undefined, metadata: undefined };
|
||||
|
||||
const params: ProjectUpdateParams = {
|
||||
...buildProjectUpdateParams(submitted, project.litellm_budget_table?.max_budget),
|
||||
team_id: submitted.team_id,
|
||||
};
|
||||
const built = buildProjectUpdateParams(submitted, project.litellm_budget_table?.max_budget);
|
||||
const storedBatchLimitKept =
|
||||
built.metadata !== undefined && !isProxyAdminRole(userRole ?? "")
|
||||
? { metadata: withStoredBatchEnqueuedTokenLimit(built.metadata, project.metadata) }
|
||||
: {};
|
||||
const params: ProjectUpdateParams = { ...built, ...storedBatchLimitKept, team_id: submitted.team_id };
|
||||
|
||||
updateMutation.mutate(
|
||||
{ projectId: project.project_id, params },
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import { useFieldArray, useWatch, type UseFormReturn } from "react-hook-form";
|
|||
import { ChevronDown, CircleAlert, Minus, Plus } from "lucide-react";
|
||||
|
||||
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
|
||||
import { isProxyAdminRole } from "@/utils/roles";
|
||||
import { ALL_TEAM_MODELS, type ProjectFormValues, type ProjectSubmitValues } from "./projectFormSchema";
|
||||
import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
|
||||
import { Team } from "@/components/key_team_helpers/key_list";
|
||||
|
|
@ -44,6 +45,7 @@ export type { ProjectFormValues } from "./projectFormSchema";
|
|||
export function ProjectBaseForm({ form, advancedOpen, onAdvancedOpenChange }: ProjectBaseFormProps) {
|
||||
const { t } = useTranslation();
|
||||
const { accessToken, userId, userRole } = useAuthorized();
|
||||
const canEditBatchLimit = isProxyAdminRole(userRole ?? "");
|
||||
const { data: teams } = useTeams();
|
||||
|
||||
const [selectedTeam, setSelectedTeam] = useState<Team | null>(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) => (
|
||||
<div
|
||||
key={field.id}
|
||||
className="mb-2 grid grid-cols-1 items-start gap-2 sm:grid-cols-2 xl:grid-cols-[minmax(0,2fr)_repeat(4,minmax(0,1fr))_auto]"
|
||||
className="mb-2 grid grid-cols-1 items-start gap-2 sm:grid-cols-2 xl:grid-cols-[minmax(0,2fr)_repeat(5,minmax(0,1fr))_auto]"
|
||||
>
|
||||
<FormField
|
||||
control={form.control}
|
||||
|
|
@ -376,6 +379,25 @@ export function ProjectBaseForm({ form, advancedOpen, onAdvancedOpenChange }: Pr
|
|||
/>
|
||||
)}
|
||||
</FormField>
|
||||
<FormField
|
||||
control={form.control}
|
||||
name={`modelLimits.${index}.batchEnqueuedTokens`}
|
||||
label={t("projects.form.modelLimits.batchEnqueuedTokens")}
|
||||
>
|
||||
{({ ref, value, onChange, ...control }) => (
|
||||
<Input
|
||||
{...control}
|
||||
ref={ref}
|
||||
type="number"
|
||||
min={1}
|
||||
step={1}
|
||||
disabled={!canEditBatchLimit}
|
||||
placeholder={t("projects.form.modelLimits.batchEnqueuedTokens")}
|
||||
value={value ?? ""}
|
||||
onChange={(event) => onChange(toOptionalNumber(event.target.value))}
|
||||
/>
|
||||
)}
|
||||
</FormField>
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
|
|
|
|||
|
|
@ -13,6 +13,11 @@ const modelLimitSchema = z.object({
|
|||
rpm: z.number().optional(),
|
||||
itpm: z.number().optional(),
|
||||
otpm: z.number().optional(),
|
||||
batchEnqueuedTokens: z
|
||||
.number()
|
||||
.int({ error: () => i18n.t("projects.form.errors.batchEnqueuedTokens") })
|
||||
.positive({ error: () => i18n.t("projects.form.errors.batchEnqueuedTokens") })
|
||||
.optional(),
|
||||
});
|
||||
|
||||
export const projectFormSchema = z
|
||||
|
|
|
|||
|
|
@ -1,5 +1,9 @@
|
|||
import { describe, it, expect } from "vitest";
|
||||
import { buildProjectCreateParams, buildProjectUpdateParams } from "./projectFormUtils";
|
||||
import {
|
||||
buildProjectCreateParams,
|
||||
buildProjectUpdateParams,
|
||||
withStoredBatchEnqueuedTokenLimit,
|
||||
} from "./projectFormUtils";
|
||||
import { ProjectFormValues } from "./ProjectBaseForm";
|
||||
|
||||
const baseValues: ProjectFormValues = {
|
||||
|
|
@ -152,6 +156,45 @@ describe("buildProjectCreateParams", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("per-model batch enqueued token limit", () => {
|
||||
it("should send the limit inside metadata next to the key-value entries", () => {
|
||||
const result = buildProjectCreateParams({
|
||||
...baseValues,
|
||||
modelLimits: [{ model: "gpt-4", batchEnqueuedTokens: 5000 }],
|
||||
metadata: [{ key: "owner", value: "platform" }],
|
||||
});
|
||||
expect(result.metadata).toEqual({ owner: "platform", model_batch_enqueued_token_limit: { "gpt-4": 5000 } });
|
||||
expect(result).not.toHaveProperty("model_batch_enqueued_token_limit");
|
||||
});
|
||||
|
||||
it("should leave the limit out of metadata when no row sets one", () => {
|
||||
const result = buildProjectUpdateParams({ ...baseValues, modelLimits: [{ model: "gpt-4", rpm: 20 }] });
|
||||
expect(result.metadata).toEqual({});
|
||||
});
|
||||
|
||||
it("should put back the stored limit that the form dropped", () => {
|
||||
const stored = { owner: "old", model_batch_enqueued_token_limit: { "gpt-4": 1000 } };
|
||||
expect(withStoredBatchEnqueuedTokenLimit({ owner: "platform" }, stored)).toStrictEqual({
|
||||
owner: "platform",
|
||||
model_batch_enqueued_token_limit: { "gpt-4": 1000 },
|
||||
});
|
||||
});
|
||||
|
||||
it("should replace a limit the form changed with the stored one", () => {
|
||||
const stored = { model_batch_enqueued_token_limit: { "gpt-4": 1000 } };
|
||||
const edited = { owner: "platform", model_batch_enqueued_token_limit: { "gpt-4": 9000 } };
|
||||
expect(withStoredBatchEnqueuedTokenLimit(edited, stored)).toStrictEqual({
|
||||
owner: "platform",
|
||||
model_batch_enqueued_token_limit: { "gpt-4": 1000 },
|
||||
});
|
||||
});
|
||||
|
||||
it("should drop a limit the form added when none is stored", () => {
|
||||
const edited = { owner: "platform", model_batch_enqueued_token_limit: { "gpt-4": 9000 } };
|
||||
expect(withStoredBatchEnqueuedTokenLimit(edited, null)).toStrictEqual({ owner: "platform" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildProjectUpdateParams clearing", () => {
|
||||
it("should send empty limit maps when every model limit row has been removed", () => {
|
||||
const result = buildProjectUpdateParams({ ...baseValues, modelLimits: [] });
|
||||
|
|
|
|||
|
|
@ -16,6 +16,30 @@ const buildModelLimitMap = (
|
|||
const buildMetadata = (entries: ProjectFormValues["metadata"]): Record<string, string> | undefined =>
|
||||
entries && Object.fromEntries(entries.flatMap((entry) => (entry.key ? [[entry.key, entry.value] as const] : [])));
|
||||
|
||||
const withBatchEnqueuedTokenLimit = (
|
||||
metadata: Record<string, string> | undefined,
|
||||
batchEnqueuedTokenLimit: Record<string, number>,
|
||||
clearsRemovedLimits: boolean,
|
||||
): Record<string, unknown> | undefined => {
|
||||
if (Object.keys(batchEnqueuedTokenLimit).length > 0) {
|
||||
return { ...metadata, model_batch_enqueued_token_limit: batchEnqueuedTokenLimit };
|
||||
}
|
||||
return clearsRemovedLimits ? { ...metadata } : metadata;
|
||||
};
|
||||
|
||||
const BATCH_ENQUEUED_TOKEN_LIMIT_KEY = "model_batch_enqueued_token_limit";
|
||||
|
||||
export const withStoredBatchEnqueuedTokenLimit = (
|
||||
metadata: Record<string, unknown>,
|
||||
storedMetadata: Record<string, unknown> | null,
|
||||
): Record<string, unknown> => {
|
||||
const withoutLimit = Object.fromEntries(
|
||||
Object.entries(metadata).filter(([key]) => key !== BATCH_ENQUEUED_TOKEN_LIMIT_KEY),
|
||||
);
|
||||
const storedLimit = storedMetadata?.[BATCH_ENQUEUED_TOKEN_LIMIT_KEY];
|
||||
return storedLimit === undefined ? withoutLimit : { ...withoutLimit, [BATCH_ENQUEUED_TOKEN_LIMIT_KEY]: storedLimit };
|
||||
};
|
||||
|
||||
const roundBudget = (value: number): number => {
|
||||
const rounded = Math.round(value * 100) / 100;
|
||||
return Number.isFinite(rounded) ? rounded : value;
|
||||
|
|
@ -27,8 +51,12 @@ const buildProjectApiParams = (values: ProjectFormValues, sendEmpty: boolean) =>
|
|||
const modelTpmLimit = buildModelLimitMap(limitEntries, (entry) => entry.tpm);
|
||||
const modelItpmLimit = buildModelLimitMap(limitEntries, (entry) => entry.itpm);
|
||||
const modelOtpmLimit = buildModelLimitMap(limitEntries, (entry) => entry.otpm);
|
||||
const metadata = buildMetadata(values.metadata);
|
||||
const keepLimits = sendEmpty && values.modelLimits !== undefined;
|
||||
const metadata = withBatchEnqueuedTokenLimit(
|
||||
buildMetadata(values.metadata),
|
||||
buildModelLimitMap(limitEntries, (entry) => entry.batchEnqueuedTokens),
|
||||
keepLimits,
|
||||
);
|
||||
const keep = (map: Record<string, number>) => keepLimits || Object.keys(map).length > 0;
|
||||
const guardrailsParam =
|
||||
values.guardrails !== undefined && (sendEmpty || values.guardrails.length > 0)
|
||||
|
|
|
|||
|
|
@ -3267,6 +3267,7 @@
|
|||
"rpm": "RPM Limit",
|
||||
"itpm": "Input TPM Limit",
|
||||
"otpm": "Output TPM Limit",
|
||||
"batchEnqueuedTokens": "Batch Enqueued Token Limit",
|
||||
"remove": "Remove model limit {{number}}",
|
||||
"add": "Add Model Limit"
|
||||
},
|
||||
|
|
@ -3284,7 +3285,8 @@
|
|||
"missingKey": "Missing key",
|
||||
"missingValue": "Missing value",
|
||||
"duplicateModel": "Duplicate model",
|
||||
"duplicateKey": "Duplicate key"
|
||||
"duplicateKey": "Duplicate key",
|
||||
"batchEnqueuedTokens": "Enter a whole number above 0"
|
||||
}
|
||||
},
|
||||
"createModal": {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,188 @@
|
|||
import json
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value
|
||||
from integration._support.process import OwnedState, owned_proxy
|
||||
from integration._support.upstream import delete_scenario, register_scenario
|
||||
from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
|
||||
from pydantic import JsonValue
|
||||
|
||||
ROWS_PER_BATCH: Final = 3
|
||||
WORDS_PER_ROW: Final = 200
|
||||
ONE_BATCH_BUT_NOT_TWO: Final = 2000
|
||||
|
||||
CUSTOM_AUTH_DOUBLE: Final = """
|
||||
import json
|
||||
import os
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
|
||||
async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth:
|
||||
if api_key == os.environ["LITELLM_MASTER_KEY"]:
|
||||
return UserAPIKeyAuth(api_key=api_key, user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
callers = json.loads(os.environ["INTEGRATION_CUSTOM_AUTH_CALLERS"])
|
||||
if api_key not in callers:
|
||||
raise Exception("unknown caller")
|
||||
return UserAPIKeyAuth(api_key=api_key, **callers[api_key])
|
||||
"""
|
||||
|
||||
|
||||
def _custom_auth_config(directory: Path) -> Path:
|
||||
(directory / "project_custom_auth.py").write_text(CUSTOM_AUTH_DOUBLE)
|
||||
base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
path: Final = directory / "project-custom-auth.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**base,
|
||||
"general_settings": {
|
||||
**object_value(base["general_settings"]),
|
||||
"custom_auth": "project_custom_auth.user_api_key_auth",
|
||||
"custom_auth_run_common_checks": True,
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def _input_lines(model_name: str) -> str:
|
||||
return "".join(
|
||||
json.dumps(
|
||||
{
|
||||
"custom_id": f"r{index}",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": model_name, "messages": [{"role": "user", "content": "enqueued " * WORDS_PER_ROW}]},
|
||||
},
|
||||
separators=(",", ":"),
|
||||
)
|
||||
+ "\n"
|
||||
for index in range(ROWS_PER_BATCH)
|
||||
)
|
||||
|
||||
|
||||
def _batch(status: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": "batch-$REQUEST_ID",
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"errors": None,
|
||||
"input_file_id": "file-in-$REQUEST_ID",
|
||||
"completion_window": "24h",
|
||||
"status": status,
|
||||
"output_file_id": None,
|
||||
"error_file_id": None,
|
||||
"created_at": 1,
|
||||
"expires_at": 1,
|
||||
"request_counts": {"total": ROWS_PER_BATCH, "completed": 0, "failed": 0},
|
||||
"metadata": None,
|
||||
}
|
||||
|
||||
|
||||
def _batch_provider(scenario: Scenario) -> str:
|
||||
model_name: Final = f"integration-{uuid.uuid4().hex}"
|
||||
handle: Final = register_scenario(
|
||||
f"project-batch-limit-{uuid.uuid4().hex[:12]}",
|
||||
RoutedResponse(
|
||||
content_type="application/x-routed",
|
||||
routes={
|
||||
"POST /files": JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
"id": "file-in-$REQUEST_ID",
|
||||
"object": "file",
|
||||
"purpose": "batch",
|
||||
"bytes": 100,
|
||||
"created_at": 1,
|
||||
"filename": "in.jsonl",
|
||||
"status": "processed",
|
||||
},
|
||||
),
|
||||
"GET /files/file-in-$REQUEST_ID/content": TextResponse(
|
||||
content_type="application/jsonl", body=_input_lines(model_name)
|
||||
),
|
||||
"POST /batches": JsonResponse(content_type="application/json", body=_batch("validating")),
|
||||
"GET /batches/batch-$REQUEST_ID": JsonResponse(
|
||||
content_type="application/json", body=_batch("cancelled")
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
scenario.cleanups.callback(delete_scenario, handle)
|
||||
scenario.model(name=model_name, api_base=handle.api_base())
|
||||
return model_name
|
||||
|
||||
|
||||
def _submit(candidate: Gateway, model_name: str, caller: str) -> httpx.Response:
|
||||
uploaded: Final = candidate.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "model": model_name},
|
||||
{"file": ("in.jsonl", _input_lines(model_name).encode(), "application/jsonl")},
|
||||
key=caller,
|
||||
)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
return candidate.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{
|
||||
"input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]),
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"model": model_name,
|
||||
},
|
||||
key=caller,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(480)
|
||||
def test_custom_auth_callers_share_the_project_per_model_enqueued_limit_until_a_batch_ends(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
limited: Final = _batch_provider(scenario)
|
||||
unlimited: Final = _batch_provider(scenario)
|
||||
team: Final = scenario.team()
|
||||
project: Final = scenario.project(
|
||||
team, metadata={"model_batch_enqueued_token_limit": {limited: ONE_BATCH_BUT_NOT_TWO}}
|
||||
)
|
||||
first_caller: Final = "sk-custom-" + uuid.uuid4().hex
|
||||
rotated_caller: Final = "sk-custom-" + uuid.uuid4().hex
|
||||
callers: Final = json.dumps(
|
||||
{
|
||||
first_caller: {"user_id": "project-batch-caller", "team_id": team, "project_id": project},
|
||||
rotated_caller: {"user_id": "project-batch-caller", "team_id": team, "project_id": project},
|
||||
}
|
||||
)
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"INTEGRATION_CUSTOM_AUTH_CALLERS": callers},
|
||||
config=_custom_auth_config(tmp_path),
|
||||
state=OwnedState.SHARED_WITH_WORKER,
|
||||
) as candidate:
|
||||
first: Final = _submit(candidate, limited, first_caller)
|
||||
assert first.status_code == 200, first.text
|
||||
|
||||
refused: Final = _submit(candidate, limited, rotated_caller)
|
||||
assert refused.status_code == 429, refused.text
|
||||
assert f"project_model: {project}:{limited}" in refused.text, refused.text
|
||||
|
||||
other_model: Final = _submit(candidate, unlimited, rotated_caller)
|
||||
assert other_model.status_code == 200, other_model.text
|
||||
|
||||
first_batch: Final = string_value(JSON_OBJECT.validate_json(first.content)["id"])
|
||||
ended: Final = candidate.request("GET", f"/v1/batches/{first_batch}", key=first_caller)
|
||||
assert ended.status_code == 200, ended.text
|
||||
assert JSON_OBJECT.validate_json(ended.content)["status"] == "cancelled", ended.text
|
||||
|
||||
after_refund: Final = _submit(candidate, limited, rotated_caller)
|
||||
assert after_refund.status_code == 200, after_refund.text
|
||||
|
|
@ -1241,9 +1241,9 @@ async def test_new_project_flag_on_missing_rpm_tpm_returns_400(monkeypatch):
|
|||
assert "rpm/tpm quota" in str(exc_info.value)
|
||||
|
||||
|
||||
def _project_update_mocks(monkeypatch, stored_metadata: dict) -> mock.MagicMock:
|
||||
def _project_update_mocks(monkeypatch, stored_metadata: dict, team_id: str | None = None) -> mock.MagicMock:
|
||||
existing_row = mock.MagicMock(
|
||||
team_id=None, budget_id=None, object_permission_id=None, metadata=stored_metadata
|
||||
team_id=team_id, budget_id=None, object_permission_id=None, metadata=stored_metadata
|
||||
)
|
||||
mock_prisma = mock.MagicMock()
|
||||
mock_prisma.jsonify_object = lambda data: data
|
||||
|
|
@ -1251,7 +1251,7 @@ def _project_update_mocks(monkeypatch, stored_metadata: dict) -> mock.MagicMock:
|
|||
mock_prisma.db.litellm_projecttable.update = mock.AsyncMock(return_value=mock.MagicMock())
|
||||
mock_prisma.writer_db = mock.MagicMock()
|
||||
mock_prisma.writer_db.litellm_projecttable.find_unique = mock.AsyncMock(
|
||||
return_value={"project_id": "project-update-test", "team_id": None}
|
||||
return_value={"project_id": "project-update-test", "team_id": team_id}
|
||||
)
|
||||
mock_prisma.writer_db.litellm_verificationtoken.count = mock.AsyncMock(return_value=0)
|
||||
|
||||
|
|
@ -2411,3 +2411,163 @@ async def test_project_writes_reject_an_unregistered_guardrail(monkeypatch, oper
|
|||
assert str(getattr(exc.value, "code", None) or getattr(exc.value, "status_code", "")) == "400"
|
||||
mock_prisma.db.litellm_projecttable.create.assert_not_awaited()
|
||||
mock_prisma.db.litellm_projecttable.update.assert_not_awaited()
|
||||
|
||||
|
||||
_TEAM_ADMIN: Final = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, api_key="hashed-team-admin-key", user_id="team-admin"
|
||||
)
|
||||
_PROXY_ADMIN: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key=MASTER_KEY, user_id="1234")
|
||||
_PROJECT_BATCH_LIMIT_REFUSAL: Final = "Only proxy admins can set model_batch_enqueued_token_limit on a project"
|
||||
|
||||
|
||||
def _stored_project_metadata() -> dict[str, object]:
|
||||
return {"owner": "platform", "model_batch_enqueued_token_limit": {"gpt-5.5": 1000}}
|
||||
|
||||
|
||||
def _team_admins_manage_projects(monkeypatch: pytest.MonkeyPatch, prisma: mock.MagicMock) -> None:
|
||||
team: Final = LiteLLM_TeamTable(
|
||||
team_id="test-team", members_with_roles=[Member(user_id=_TEAM_ADMIN.user_id, role="admin")]
|
||||
)
|
||||
prisma.db.litellm_teamtable.find_unique = mock.AsyncMock(return_value=team)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"team_admin_editable_team_fields": ["projects"]})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_admin_cannot_put_a_model_batch_enqueued_token_limit_on_a_new_project(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
prisma: Final = mock.MagicMock()
|
||||
_team_admins_manage_projects(monkeypatch, prisma)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", prisma)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None)
|
||||
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await new_project(
|
||||
data=NewProjectRequest(
|
||||
team_id="test-team", metadata={"model_batch_enqueued_token_limit": {"gpt-5.5": 1000}}
|
||||
),
|
||||
http_request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=_TEAM_ADMIN,
|
||||
)
|
||||
|
||||
assert exc.value.code == "403"
|
||||
assert _PROJECT_BATCH_LIMIT_REFUSAL in exc.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"requested_metadata",
|
||||
[{"owner": "platform", "model_batch_enqueued_token_limit": {"gpt-5.5": 999_999}}, {"owner": "platform"}],
|
||||
ids=["raises the limit", "drops the limit"],
|
||||
)
|
||||
async def test_team_admin_cannot_change_a_project_model_batch_enqueued_token_limit(
|
||||
monkeypatch: pytest.MonkeyPatch, requested_metadata: dict[str, object]
|
||||
) -> None:
|
||||
mock_prisma: Final = _project_update_mocks(monkeypatch, _stored_project_metadata(), team_id="test-team")
|
||||
_team_admins_manage_projects(monkeypatch, mock_prisma)
|
||||
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await update_project(
|
||||
data=UpdateProjectRequest(project_id=f"project-{uuid.uuid4()}", metadata=requested_metadata),
|
||||
http_request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=_TEAM_ADMIN,
|
||||
)
|
||||
|
||||
assert exc.value.code == "403"
|
||||
assert _PROJECT_BATCH_LIMIT_REFUSAL in exc.value.message
|
||||
mock_prisma.db.litellm_projecttable.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("caller", "requested_limit"),
|
||||
[(_TEAM_ADMIN, {"gpt-5.5": 1000}), (_PROXY_ADMIN, {"gpt-5.5": 5000})],
|
||||
ids=["team admin resends the stored limit", "proxy admin changes the limit"],
|
||||
)
|
||||
async def test_project_update_keeps_an_allowed_model_batch_enqueued_token_limit(
|
||||
monkeypatch: pytest.MonkeyPatch, caller: UserAPIKeyAuth, requested_limit: dict[str, int]
|
||||
) -> None:
|
||||
mock_prisma: Final = _project_update_mocks(monkeypatch, _stored_project_metadata(), team_id="test-team")
|
||||
_team_admins_manage_projects(monkeypatch, mock_prisma)
|
||||
|
||||
await update_project(
|
||||
data=UpdateProjectRequest(
|
||||
project_id=f"project-{uuid.uuid4()}",
|
||||
metadata={"owner": "renamed", "model_batch_enqueued_token_limit": requested_limit},
|
||||
),
|
||||
http_request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=caller,
|
||||
)
|
||||
|
||||
written_metadata: Final = _written_project_data(mock_prisma)["metadata"]
|
||||
assert written_metadata["owner"] == "renamed"
|
||||
assert written_metadata["model_batch_enqueued_token_limit"] == requested_limit
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_admin_can_write_batch_keys_a_project_does_not_enforce(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
mock_prisma: Final = _project_update_mocks(monkeypatch, _stored_project_metadata(), team_id="test-team")
|
||||
_team_admins_manage_projects(monkeypatch, mock_prisma)
|
||||
requested_metadata: Final = {
|
||||
**_stored_project_metadata(),
|
||||
"batch_enqueued_token_limit": 50,
|
||||
"max_batch_file_records": 10,
|
||||
"allowed_s3_buckets": ["team-bucket"],
|
||||
}
|
||||
|
||||
await update_project(
|
||||
data=UpdateProjectRequest(project_id=f"project-{uuid.uuid4()}", metadata=requested_metadata),
|
||||
http_request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=_TEAM_ADMIN,
|
||||
)
|
||||
|
||||
assert _written_project_data(mock_prisma)["metadata"] == requested_metadata
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"requested_limit",
|
||||
[{"gpt-5.5": 1000, "gpt-5.4": 0}, {"gpt-5.5": "lots"}, [1000], {"gpt-5.5": True}, {"gpt-5.5": "2000"}],
|
||||
ids=["one model at zero", "a non-numeric limit", "not a model map", "a boolean limit", "a numeric string limit"],
|
||||
)
|
||||
async def test_project_update_rejects_a_model_batch_enqueued_token_limit_it_cannot_enforce(
|
||||
monkeypatch: pytest.MonkeyPatch, requested_limit: object
|
||||
) -> None:
|
||||
mock_prisma: Final = _project_update_mocks(monkeypatch, _stored_project_metadata(), team_id="test-team")
|
||||
_team_admins_manage_projects(monkeypatch, mock_prisma)
|
||||
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await update_project(
|
||||
data=UpdateProjectRequest(
|
||||
project_id=f"project-{uuid.uuid4()}", metadata={"model_batch_enqueued_token_limit": requested_limit}
|
||||
),
|
||||
http_request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=_PROXY_ADMIN,
|
||||
)
|
||||
|
||||
assert exc.value.code == "400"
|
||||
assert "model_batch_enqueued_token_limit must map each model name to a positive integer" in exc.value.message
|
||||
mock_prisma.db.litellm_projecttable.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"top_level_fields",
|
||||
[{"model_tpm_limit": {"gpt-5.5": 2000}}, {"tags": ["batch"]}],
|
||||
ids=["a per-model limit", "a premium metadata field"],
|
||||
)
|
||||
@pytest.mark.parametrize("caller", [_TEAM_ADMIN, _PROXY_ADMIN], ids=["team admin", "proxy admin"])
|
||||
async def test_project_update_with_only_top_level_fields_keeps_the_stored_metadata(
|
||||
monkeypatch: pytest.MonkeyPatch, top_level_fields: dict[str, object], caller: UserAPIKeyAuth
|
||||
) -> None:
|
||||
mock_prisma: Final = _project_update_mocks(monkeypatch, _stored_project_metadata(), team_id="test-team")
|
||||
_team_admins_manage_projects(monkeypatch, mock_prisma)
|
||||
|
||||
await update_project(
|
||||
data=UpdateProjectRequest(project_id=f"project-{uuid.uuid4()}", **top_level_fields),
|
||||
http_request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=caller,
|
||||
)
|
||||
|
||||
assert _written_project_data(mock_prisma)["metadata"] == {**_stored_project_metadata(), **top_level_fields}
|
||||
|
|
|
|||
|
|
@ -10,14 +10,17 @@ import base64
|
|||
import logging
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
|
||||
from litellm.constants import BATCH_ENQUEUED_TOKEN_TTL_SECONDS
|
||||
from litellm.litellm_core_utils.model_latest_aliases import get_model_latest_aliases
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.hooks.batch_enqueued_tokens import (
|
||||
BatchEnqueuedTokenOverLimit,
|
||||
|
|
@ -26,7 +29,10 @@ from litellm.proxy.hooks.batch_enqueued_tokens import (
|
|||
BatchEnqueuedTokenStore,
|
||||
batch_response_view,
|
||||
canonical_provider_batch_id,
|
||||
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.utils import InternalUsageCache
|
||||
|
||||
|
|
@ -69,6 +75,156 @@ def test_scope_resolution_skips_team_scope_without_team_id():
|
|||
assert resolve_batch_enqueued_token_scopes(user) == ()
|
||||
|
||||
|
||||
def _project_user(model_limits: object, project_id: str | None = "project-1") -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
api_key="hashed-jwt-abc",
|
||||
team_id="shared-team",
|
||||
project_id=project_id,
|
||||
project_metadata={"model_batch_enqueued_token_limit": model_limits},
|
||||
)
|
||||
|
||||
|
||||
def test_project_model_scopes_charge_each_limited_row_model_only_its_own_tokens():
|
||||
scopes: Final = resolve_project_model_batch_enqueued_token_scopes(
|
||||
_project_user({"gpt-5.5": 100, "claude-opus-5-5": 50}),
|
||||
row_tokens_by_model={"gpt-5.5": 80, "unlimited-model": 999},
|
||||
routing_model="unlimited-model",
|
||||
batch_tokens=1079,
|
||||
model_group_of=lambda model: model,
|
||||
)
|
||||
assert scopes == (BatchEnqueuedTokenScope(key="project_model", value="project-1:gpt-5.5", limit=100, tokens=80),)
|
||||
|
||||
|
||||
def test_project_model_scopes_charge_the_routing_model_the_whole_batch_whatever_rows_name():
|
||||
scopes: Final = resolve_project_model_batch_enqueued_token_scopes(
|
||||
_project_user({"capped-deployment": 100}),
|
||||
row_tokens_by_model={"unlisted-model": 400},
|
||||
routing_model="capped-deployment",
|
||||
batch_tokens=500,
|
||||
model_group_of=lambda model: model,
|
||||
)
|
||||
assert scopes == (
|
||||
BatchEnqueuedTokenScope(key="project_model", value="project-1:capped-deployment", limit=100, tokens=500),
|
||||
)
|
||||
|
||||
|
||||
def test_project_model_scopes_charge_rows_of_the_routing_group_once():
|
||||
groups: Final = MappingProxyType({"mini-alias": "gpt-mini", "gpt-mini": "gpt-mini"})
|
||||
scopes: Final = resolve_project_model_batch_enqueued_token_scopes(
|
||||
_project_user({"gpt-mini": 100}),
|
||||
row_tokens_by_model={"mini-alias": 30},
|
||||
routing_model="gpt-mini",
|
||||
batch_tokens=30,
|
||||
model_group_of=groups.__getitem__,
|
||||
)
|
||||
assert scopes == (BatchEnqueuedTokenScope(key="project_model", value="project-1:gpt-mini", limit=100, tokens=30),)
|
||||
|
||||
|
||||
def test_project_model_scopes_share_one_counter_per_model_group_under_the_tightest_limit():
|
||||
groups: Final = MappingProxyType(
|
||||
{"gpt-latest": "gpt-5.5", "gpt-5.5": "gpt-5.5", "mini": "gpt-mini", "gpt-mini": "gpt-mini"}
|
||||
)
|
||||
scopes: Final = resolve_project_model_batch_enqueued_token_scopes(
|
||||
_project_user({"gpt-latest": 100, "gpt-5.5": 70, "gpt-mini": 40}),
|
||||
row_tokens_by_model={"gpt-latest": 30, "gpt-5.5": 20, "mini": 5},
|
||||
routing_model=None,
|
||||
batch_tokens=55,
|
||||
model_group_of=groups.__getitem__,
|
||||
)
|
||||
assert scopes == (
|
||||
BatchEnqueuedTokenScope(key="project_model", value="project-1:gpt-5.5", limit=70, tokens=50),
|
||||
BatchEnqueuedTokenScope(key="project_model", value="project-1:gpt-mini", limit=40, tokens=5),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("named_model", ["gpt-latest", "gpt-5.5"], ids=["only an alias", "only the group"])
|
||||
def test_project_model_scope_takes_its_groups_tightest_limit_whichever_name_the_batch_uses(named_model: str) -> None:
|
||||
groups: Final = MappingProxyType({"gpt-latest": "gpt-5.5", "gpt-pinned": "gpt-5.5", "gpt-5.5": "gpt-5.5"})
|
||||
scopes: Final = resolve_project_model_batch_enqueued_token_scopes(
|
||||
_project_user({"gpt-latest": 100, "gpt-5.5": 70, "gpt-pinned": 60}),
|
||||
row_tokens_by_model={named_model: 30},
|
||||
routing_model=named_model,
|
||||
batch_tokens=30,
|
||||
model_group_of=groups.__getitem__,
|
||||
)
|
||||
assert scopes == (BatchEnqueuedTokenScope(key="project_model", value="project-1:gpt-5.5", limit=60, tokens=30),)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"user",
|
||||
[
|
||||
_project_user(100),
|
||||
_project_user({"gpt-5.5": 0}),
|
||||
_project_user({"gpt-5.5": "lots"}),
|
||||
_project_user(["gpt-5.5"]),
|
||||
_project_user({"gpt-5.5": True}),
|
||||
_project_user({"gpt-5.5": "100"}),
|
||||
_project_user({"gpt-5.5": 100.0}),
|
||||
_project_user({"gpt-5.5": 100}, project_id=None),
|
||||
UserAPIKeyAuth(api_key="k", project_id="project-1"),
|
||||
],
|
||||
ids=["scalar", "zero", "non-numeric", "list", "boolean", "numeric string", "float", "no project id", "no metadata"],
|
||||
)
|
||||
def test_project_model_scopes_ignore_unusable_limits(user):
|
||||
assert not project_has_model_batch_enqueued_token_limits(user)
|
||||
assert (
|
||||
resolve_project_model_batch_enqueued_token_scopes(
|
||||
user, row_tokens_by_model={"gpt-5.5": 10}, routing_model="gpt-5.5", batch_tokens=10, model_group_of=str
|
||||
)
|
||||
== ()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("models", "expected"),
|
||||
[
|
||||
(("gpt-5.5",), True),
|
||||
(("gpt-latest",), True),
|
||||
(("unlisted-model",), False),
|
||||
(("gpt-5.5", "unlisted-model"), False),
|
||||
],
|
||||
ids=["limited model", "alias of a limited group", "unlisted model", "one of the models is unlisted"],
|
||||
)
|
||||
def test_project_limits_every_model_only_when_each_model_has_a_limit(models, expected):
|
||||
groups: Final = MappingProxyType({"gpt-latest": "gpt-5.5"})
|
||||
assert (
|
||||
project_limits_every_model(
|
||||
_project_user({"gpt-5.5": 100}), models, model_group_of=lambda name: groups.get(name, name)
|
||||
)
|
||||
is expected
|
||||
)
|
||||
|
||||
|
||||
def test_a_latest_alias_is_charged_under_the_project_limit_of_its_target():
|
||||
alias, target = next(iter(get_model_latest_aliases()["openai"].items()))
|
||||
user: Final = _project_user({target: 100})
|
||||
|
||||
assert project_limits_every_model(user, (alias,), model_group_of=str)
|
||||
assert resolve_project_model_batch_enqueued_token_scopes(
|
||||
user, row_tokens_by_model={alias: 40}, routing_model=alias, batch_tokens=40, model_group_of=str
|
||||
) == (BatchEnqueuedTokenScope(key="project_model", value=f"project-1:{target}", limit=100, tokens=40),)
|
||||
|
||||
|
||||
def test_a_dated_snapshot_row_is_charged_under_the_project_limit_of_its_base_model():
|
||||
user: Final = _project_user({"gpt-5.5": 100, "gpt-5.5-2026-04-23": 60})
|
||||
|
||||
assert project_limits_every_model(_project_user({"gpt-5.5": 100}), ("gpt-5.5-2026-04-23",), model_group_of=str)
|
||||
assert resolve_project_model_batch_enqueued_token_scopes(
|
||||
user,
|
||||
row_tokens_by_model={"gpt-5.5-2026-04-23": 40},
|
||||
routing_model="gpt-5.4-nano",
|
||||
batch_tokens=40,
|
||||
model_group_of=str,
|
||||
) == (
|
||||
BatchEnqueuedTokenScope(key="project_model", value="project-1:gpt-5.5", limit=100, tokens=40),
|
||||
BatchEnqueuedTokenScope(key="project_model", value="project-1:gpt-5.5-2026-04-23", limit=60, tokens=40),
|
||||
)
|
||||
|
||||
|
||||
def test_project_model_limits_are_detected_before_the_batch_file_is_read():
|
||||
assert project_has_model_batch_enqueued_token_limits(_project_user({"gpt-5.5": 100}))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reserve_rejects_once_allowance_is_exhausted():
|
||||
store = _in_memory_store()
|
||||
|
|
@ -216,6 +372,89 @@ async def test_redis_reserve_issues_single_key_calls_and_rolls_back_on_over_limi
|
|||
assert all(len(keys) == 1 for _, keys in fake.script_calls)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_project_model_scope_is_charged_its_own_tokens_and_refunded_after_a_round_trip():
|
||||
fake: Final = _SingleKeyRedisFake()
|
||||
store: Final = BatchEnqueuedTokenStore(
|
||||
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
|
||||
)
|
||||
key_scope: Final = _scope(limit=1000, key="api_key")
|
||||
model_scope: Final = BatchEnqueuedTokenScope(
|
||||
key="project_model", value=f"project-{uuid.uuid4().hex}:gpt-5.5", limit=40, tokens=30
|
||||
)
|
||||
key_counter: Final = f"batch_enqueued_tokens:api_key:{key_scope.value}"
|
||||
model_counter: Final = f"batch_enqueued_tokens:project_model:{model_scope.value}"
|
||||
|
||||
reservation: Final = await store.reserve(tokens=100, scopes=(key_scope, model_scope))
|
||||
assert isinstance(reservation, BatchEnqueuedTokenReservation)
|
||||
assert fake.counters == {key_counter: 100, model_counter: 30}
|
||||
|
||||
over: Final = await store.reserve(
|
||||
tokens=100,
|
||||
scopes=(key_scope, BatchEnqueuedTokenScope(key="project_model", value=model_scope.value, limit=40, tokens=20)),
|
||||
)
|
||||
assert isinstance(over, BatchEnqueuedTokenOverLimit)
|
||||
assert over.enqueued == 30
|
||||
assert fake.counters == {key_counter: 100, model_counter: 30}
|
||||
|
||||
await store.save_reservation("batch_project_model", reservation)
|
||||
popped: Final = await store.pop_reservation("batch_project_model")
|
||||
assert popped == reservation
|
||||
await store.refund(popped)
|
||||
assert not fake.counters
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ScopeBeforeProjectScopes:
|
||||
key: Literal["api_key", "team"]
|
||||
value: str
|
||||
limit: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ReservationBeforeProjectScopes:
|
||||
tokens: int
|
||||
scopes: tuple[_ScopeBeforeProjectScopes, ...]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_pod_that_predates_project_scopes_still_reads_the_key_and_team_scopes_of_a_saved_record():
|
||||
fake: Final = _SingleKeyRedisFake()
|
||||
store: Final = BatchEnqueuedTokenStore(
|
||||
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60))
|
||||
)
|
||||
key_scope: Final = _scope(limit=1000, key="api_key")
|
||||
model_scope: Final = BatchEnqueuedTokenScope(key="project_model", value="project-1:gpt-5.5", limit=40, tokens=30)
|
||||
reservation: Final = await store.reserve(tokens=100, scopes=(key_scope, model_scope))
|
||||
assert isinstance(reservation, BatchEnqueuedTokenReservation)
|
||||
|
||||
await store.save_reservation("batch_mixed_versions", reservation)
|
||||
|
||||
(record,) = fake.records.values()
|
||||
assert TypeAdapter(_ReservationBeforeProjectScopes).validate_json(record) == _ReservationBeforeProjectScopes(
|
||||
tokens=100, scopes=(_ScopeBeforeProjectScopes(key="api_key", value=key_scope.value, limit=1000),)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_in_memory_project_model_scope_is_charged_and_refunded_its_own_tokens():
|
||||
store: Final = _in_memory_store()
|
||||
key_scope: Final = _scope(limit=1000, key="api_key")
|
||||
model_scope: Final = BatchEnqueuedTokenScope(
|
||||
key="project_model", value=f"project-{uuid.uuid4().hex}:gpt-5.5", limit=40, tokens=30
|
||||
)
|
||||
reservation: Final = await store.reserve(tokens=100, scopes=(key_scope, model_scope))
|
||||
assert isinstance(reservation, BatchEnqueuedTokenReservation)
|
||||
blocked: Final = BatchEnqueuedTokenScope(key="project_model", value=model_scope.value, limit=40, tokens=20)
|
||||
assert await store.reserve(tokens=1, scopes=(blocked,)) == BatchEnqueuedTokenOverLimit(scope=blocked, enqueued=30)
|
||||
|
||||
await store.refund(reservation)
|
||||
full: Final = BatchEnqueuedTokenScope(key="project_model", value=model_scope.value, limit=40, tokens=40)
|
||||
assert isinstance(await store.reserve(tokens=1, scopes=(full,)), BatchEnqueuedTokenReservation)
|
||||
assert isinstance(await store.reserve(tokens=900, scopes=(key_scope,)), BatchEnqueuedTokenReservation)
|
||||
assert isinstance(await store.reserve(tokens=101, scopes=(key_scope,)), BatchEnqueuedTokenOverLimit)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_partial_redis_reserve_failure_rolls_back_and_grants_in_memory():
|
||||
key_scope = _scope(limit=100, key="api_key")
|
||||
|
|
|
|||
|
|
@ -7,6 +7,8 @@ VERIA-39 regression tests:
|
|||
models the caller is not authorized to use.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -2101,7 +2103,10 @@ def test_estimate_entry_output_tokens_multiplies_candidate_count(body_extra, exp
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _enqueued_rate_limiter():
|
||||
def _enqueued_rate_limiter(
|
||||
model_group_resolver: Callable[[str], str | None] = lambda _model: None,
|
||||
deployment_model_name: Callable[[str], str | None] = lambda _model: None,
|
||||
):
|
||||
from litellm import DualCache
|
||||
from litellm.proxy.hooks.batch_rate_limiter import PROXY_BatchRateLimiter
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
|
|
@ -2111,10 +2116,13 @@ def _enqueued_rate_limiter():
|
|||
|
||||
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)
|
||||
parallel_request_limiter: Final = PROXY_MaxParallelRequestsHandler_v3(
|
||||
internal_usage_cache=internal_usage_cache, model_group_resolver=model_group_resolver
|
||||
)
|
||||
rate_limiter = PROXY_BatchRateLimiter(
|
||||
internal_usage_cache=internal_usage_cache,
|
||||
parallel_request_limiter=parallel_request_limiter,
|
||||
deployment_model_name=deployment_model_name,
|
||||
)
|
||||
return rate_limiter, local_cache
|
||||
|
||||
|
|
@ -2126,9 +2134,9 @@ _ENQUEUED_BATCH_FILE_CONTENT = (
|
|||
)
|
||||
|
||||
|
||||
def _enqueued_batch_patches():
|
||||
def _enqueued_batch_patches(file_content: bytes = _ENQUEUED_BATCH_FILE_CONTENT):
|
||||
mock_content = MagicMock()
|
||||
mock_content.content = _ENQUEUED_BATCH_FILE_CONTENT
|
||||
mock_content.content = file_content
|
||||
afile_content_mock = AsyncMock(return_value=mock_content)
|
||||
return afile_content_mock, (
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||||
|
|
@ -2261,6 +2269,307 @@ async def test_enqueued_team_limit_applies_to_batch_submission():
|
|||
assert "Batch enqueued token limit exceeded for team: team-enqueued-batch" in str(exc.value.detail)
|
||||
|
||||
|
||||
_MODEL_LESS_BATCH_FILE_CONTENT = (
|
||||
b'{"body": {"max_tokens": 40, "messages": [{"role": "user", "content": "hi"}]}}\n'
|
||||
b'{"body": {"max_tokens": 40, "messages": [{"role": "user", "content": "hi"}]}}\n'
|
||||
)
|
||||
_ALIASED_BATCH_FILE_CONTENT = _ENQUEUED_BATCH_FILE_CONTENT.replace(b"gpt-4o-mini", b"mini-alias")
|
||||
_UNLISTED_MODEL_BATCH_FILE_CONTENT = _ENQUEUED_BATCH_FILE_CONTENT.replace(b"gpt-4o-mini", b"unlisted-model")
|
||||
_MIXED_MODEL_BATCH_FILE_CONTENT = (
|
||||
_ENQUEUED_BATCH_FILE_CONTENT
|
||||
+ b'{"body": {"model": "gpt-5.5", "max_tokens": 40, "messages": [{"role": "user", "content": "hello there"}]}}\n'
|
||||
)
|
||||
|
||||
|
||||
def _project_model_limited_user(model_limits: Mapping[str, int]) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
api_key="hashed-jwt-project-caller",
|
||||
models=["*"],
|
||||
team_id="shared-team",
|
||||
project_id="project-enqueued-batch",
|
||||
project_metadata={"model_batch_enqueued_token_limit": model_limits},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"file_content",
|
||||
[
|
||||
_ENQUEUED_BATCH_FILE_CONTENT,
|
||||
_MODEL_LESS_BATCH_FILE_CONTENT,
|
||||
_ALIASED_BATCH_FILE_CONTENT,
|
||||
_UNLISTED_MODEL_BATCH_FILE_CONTENT,
|
||||
],
|
||||
ids=[
|
||||
"rows name the model",
|
||||
"rows name no model",
|
||||
"rows name an alias of the model",
|
||||
"rows name an unlisted model",
|
||||
],
|
||||
)
|
||||
async def test_project_model_enqueued_limit_rejects_a_batch_over_that_models_limit(file_content):
|
||||
rate_limiter, local_cache = _enqueued_rate_limiter({"mini-alias": "gpt-4o-mini"}.get)
|
||||
afile_content_mock, patches = _enqueued_batch_patches(file_content)
|
||||
|
||||
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=_project_model_limited_user({"gpt-4o-mini": 10}),
|
||||
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 project_model: project-enqueued-batch:gpt-4o-mini" in str(
|
||||
exc.value.detail
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_project_model_enqueued_limit_reserves_only_the_limited_models_rows():
|
||||
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(_MIXED_MODEL_BATCH_FILE_CONTENT)
|
||||
|
||||
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
|
||||
data: Final = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
|
||||
await rate_limiter.async_pre_call_hook(
|
||||
user_api_key_dict=_project_model_limited_user({"gpt-5.5": 100000}),
|
||||
cache=local_cache,
|
||||
data=data,
|
||||
call_type="acreate_batch",
|
||||
)
|
||||
|
||||
stash: Final = get_request_stash()
|
||||
assert stash is not None and stash.batch_enqueued_reservation is not None
|
||||
(scope,) = stash.batch_enqueued_reservation.scopes
|
||||
assert scope.value == "project-enqueued-batch:gpt-5.5"
|
||||
assert scope.tokens is not None
|
||||
assert 0 < scope.tokens < data["_batch_token_count"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_project_model_enqueued_limit_for_another_model_leaves_the_batch_on_per_minute_limits():
|
||||
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()
|
||||
|
||||
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
|
||||
data: Final = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
|
||||
result: Final = await rate_limiter.async_pre_call_hook(
|
||||
user_api_key_dict=_project_model_limited_user({"gpt-5.5": 10}),
|
||||
cache=local_cache,
|
||||
data=data,
|
||||
call_type="acreate_batch",
|
||||
)
|
||||
|
||||
assert result is data
|
||||
assert data["_batch_token_count"] > 10
|
||||
stash: Final = get_request_stash()
|
||||
assert stash is None or stash.batch_enqueued_reservation is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_minute_limits_still_apply_when_only_a_rows_model_has_a_project_allowance():
|
||||
from litellm.proxy.hooks.batch_enqueued_tokens import BatchEnqueuedTokenScope
|
||||
|
||||
rate_limiter, local_cache = _enqueued_rate_limiter()
|
||||
afile_content_mock, patches = _enqueued_batch_patches(_MIXED_MODEL_BATCH_FILE_CONTENT)
|
||||
allowance: Final = 1000000
|
||||
|
||||
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
|
||||
per_minute_capped: Final = _project_model_limited_user({"gpt-5.5": allowance}).model_copy(
|
||||
update={"rpm_limit": 1, "tpm_limit": 10}
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await rate_limiter.async_pre_call_hook(
|
||||
user_api_key_dict=per_minute_capped,
|
||||
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" not in str(exc.value.detail)
|
||||
whole_allowance: Final = BatchEnqueuedTokenScope(
|
||||
key="project_model", value="project-enqueued-batch:gpt-5.5", limit=allowance, tokens=allowance
|
||||
)
|
||||
refilled: Final = await rate_limiter.parallel_request_limiter.batch_enqueued_token_store.reserve(
|
||||
tokens=allowance, scopes=(whole_allowance,)
|
||||
)
|
||||
assert refilled.scopes == (whole_allowance,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_project_allowance_for_the_routing_model_replaces_per_minute_limits():
|
||||
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()
|
||||
|
||||
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
|
||||
per_minute_capped: Final = _project_model_limited_user({"gpt-4o-mini": 1000000}).model_copy(
|
||||
update={"rpm_limit": 1, "tpm_limit": 10}
|
||||
)
|
||||
data: Final = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
|
||||
result: Final = await rate_limiter.async_pre_call_hook(
|
||||
user_api_key_dict=per_minute_capped,
|
||||
cache=local_cache,
|
||||
data=data,
|
||||
call_type="acreate_batch",
|
||||
)
|
||||
|
||||
assert result is data
|
||||
stash: Final = get_request_stash()
|
||||
assert stash is not None and stash.batch_enqueued_reservation is not None
|
||||
assert stash.batch_enqueued_reservation.tokens == data["_batch_token_count"] > 10
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_routing_model_allowance_keeps_project_input_limits_on_the_other_row_models():
|
||||
from litellm.proxy.hooks.batch_enqueued_tokens import BatchEnqueuedTokenScope
|
||||
|
||||
rate_limiter, local_cache = _enqueued_rate_limiter()
|
||||
afile_content_mock, patches = _enqueued_batch_patches(_MIXED_MODEL_BATCH_FILE_CONTENT)
|
||||
allowance: Final = 1000000
|
||||
user: Final = _project_model_limited_user({"gpt-4o-mini": allowance}).model_copy(
|
||||
update={
|
||||
"project_metadata": {
|
||||
"model_batch_enqueued_token_limit": {"gpt-4o-mini": allowance},
|
||||
"model_itpm_limit": {"gpt-5.5": 1},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
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=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 "model_per_project_itpm" in str(exc.value.detail)
|
||||
whole_allowance: Final = BatchEnqueuedTokenScope(
|
||||
key="project_model", value="project-enqueued-batch:gpt-4o-mini", limit=allowance, tokens=allowance
|
||||
)
|
||||
refilled: Final = await rate_limiter.parallel_request_limiter.batch_enqueued_token_store.reserve(
|
||||
tokens=allowance, scopes=(whole_allowance,)
|
||||
)
|
||||
assert refilled.scopes == (whole_allowance,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_daily_token_limit_still_applies_alongside_a_project_model_reservation():
|
||||
rate_limiter, local_cache = _enqueued_rate_limiter()
|
||||
afile_content_mock, patches = _enqueued_batch_patches()
|
||||
|
||||
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
|
||||
daily_capped: Final = _project_model_limited_user({"gpt-4o-mini": 1000000}).model_copy(update={"tpd_limit": 10})
|
||||
rejected_data: Final = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await rate_limiter.async_pre_call_hook(
|
||||
user_api_key_dict=daily_capped,
|
||||
cache=local_cache,
|
||||
data=rejected_data,
|
||||
call_type="acreate_batch",
|
||||
)
|
||||
assert exc.value.status_code == 429
|
||||
assert "Batch enqueued token limit" not in str(exc.value.detail)
|
||||
|
||||
exact_fit: Final = _project_model_limited_user({"gpt-4o-mini": rejected_data["_batch_token_count"]})
|
||||
accepted_data: Final = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
|
||||
result: Final = await rate_limiter.async_pre_call_hook(
|
||||
user_api_key_dict=exact_fit,
|
||||
cache=local_cache,
|
||||
data=accepted_data,
|
||||
call_type="acreate_batch",
|
||||
)
|
||||
|
||||
assert result is accepted_data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"opted_in_user",
|
||||
[
|
||||
UserAPIKeyAuth(
|
||||
api_key="sk-enqueued-daily", models=["*"], tpd_limit=10, metadata={"batch_enqueued_token_limit": 1000000}
|
||||
),
|
||||
UserAPIKeyAuth(
|
||||
api_key="sk-team-enqueued-daily",
|
||||
models=["*"],
|
||||
team_id="team-enqueued-daily",
|
||||
team_tpd_limit=10,
|
||||
team_metadata={"batch_enqueued_token_limit": 1000000},
|
||||
),
|
||||
],
|
||||
ids=["key", "team"],
|
||||
)
|
||||
async def test_key_or_team_enqueued_limit_alone_governs_a_batch_over_its_daily_cap(opted_in_user):
|
||||
rate_limiter, local_cache = _enqueued_rate_limiter()
|
||||
afile_content_mock, patches = _enqueued_batch_patches()
|
||||
|
||||
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
|
||||
data: Final = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
|
||||
result: Final = 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
|
||||
assert data["_batch_token_count"] > 10
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_model_bound_upload_still_charges_the_model_its_rows_name():
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
|
||||
|
||||
rate_limiter, local_cache = _enqueued_rate_limiter()
|
||||
afile_content_mock, patches = _enqueued_batch_patches()
|
||||
|
||||
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=_project_model_limited_user({"gpt-4o-mini-priority": 1000000, "gpt-4o-mini": 10}),
|
||||
cache=local_cache,
|
||||
data={"input_file_id": encode_file_id_with_model("file-abc123", "gpt-4o-mini-priority")},
|
||||
call_type="acreate_batch",
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 429
|
||||
assert "Batch enqueued token limit exceeded for project_model: project-enqueued-batch:gpt-4o-mini." in str(
|
||||
exc.value.detail
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_project_model_enqueued_limit_applies_to_a_batch_that_names_a_deployment_id():
|
||||
rate_limiter, local_cache = _enqueued_rate_limiter(deployment_model_name={"deployment-1": "gpt-4o-mini"}.get)
|
||||
afile_content_mock, patches = _enqueued_batch_patches(_MODEL_LESS_BATCH_FILE_CONTENT)
|
||||
|
||||
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=_project_model_limited_user({"gpt-4o-mini": 10}),
|
||||
cache=local_cache,
|
||||
data={"input_file_id": "file-abc123", "model": "deployment-1"},
|
||||
call_type="acreate_batch",
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 429
|
||||
assert "Batch enqueued token limit exceeded for project_model: project-enqueued-batch:gpt-4o-mini" 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()
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ batch under a per-minute RPM/TPM budget. Scopes that configure `tpd_limit`
|
|||
are charged against a 24h token window instead of their minute counters.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
|
|
@ -555,18 +556,13 @@ def _serve_file_to_the_deployment_key(rows: Sequence[_BatchLine]) -> Callable[[h
|
|||
return respond
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"query_string, routing_headers",
|
||||
[(b"model=routed-batch-model", ()), (b"", ((b"x-litellm-model", b"routed-batch-model"),))],
|
||||
ids=["query", "header"],
|
||||
)
|
||||
async def test_model_routed_batch_is_counted_with_the_deployment_key_not_the_request_key(
|
||||
async def _model_routed_batch_data(
|
||||
openai_files: respx.MockRouter,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
query_string: bytes,
|
||||
routing_headers: tuple[tuple[bytes, bytes], ...],
|
||||
):
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> tuple[dict[str, object], respx.Route]:
|
||||
from starlette.requests import Request
|
||||
|
||||
from litellm import Router
|
||||
|
|
@ -593,7 +589,6 @@ async def test_model_routed_batch_is_counted_with_the_deployment_key_not_the_req
|
|||
file_route: Final = openai_files.get(_ROUTED_FILE_URL).mock(
|
||||
side_effect=_serve_file_to_the_deployment_key(_batch_rows(("Hello " * 50,)))
|
||||
)
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(api_key=hash_token("sk-batch"), tpm_limit=1)
|
||||
request: Final = Request(
|
||||
{
|
||||
"type": "http",
|
||||
|
|
@ -620,6 +615,28 @@ async def test_model_routed_batch_is_counted_with_the_deployment_key_not_the_req
|
|||
proxy_config=MagicMock(),
|
||||
general_settings={},
|
||||
)
|
||||
return data, file_route
|
||||
|
||||
|
||||
_MODEL_ROUTES: Final = pytest.mark.parametrize(
|
||||
"query_string, routing_headers",
|
||||
[(b"model=routed-batch-model", ()), (b"", ((b"x-litellm-model", b"routed-batch-model"),))],
|
||||
ids=["query", "header"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@_MODEL_ROUTES
|
||||
async def test_model_routed_batch_is_counted_with_the_deployment_key_not_the_request_key(
|
||||
openai_files: respx.MockRouter,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
query_string: bytes,
|
||||
routing_headers: tuple[tuple[bytes, bytes], ...],
|
||||
):
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(api_key=hash_token("sk-batch"), tpm_limit=1)
|
||||
data, file_route = await _model_routed_batch_data(
|
||||
openai_files, monkeypatch, query_string, routing_headers, user_api_key_dict
|
||||
)
|
||||
_, _, batch_limiter = _make_limiters()
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
|
|
@ -631,8 +648,37 @@ async def test_model_routed_batch_is_counted_with_the_deployment_key_not_the_req
|
|||
assert file_route.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@_MODEL_ROUTES
|
||||
async def test_model_routed_batch_is_charged_to_the_project_cap_of_the_model_it_routes_to(
|
||||
openai_files: respx.MockRouter,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
query_string: bytes,
|
||||
routing_headers: tuple[tuple[bytes, bytes], ...],
|
||||
):
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-batch"),
|
||||
project_id="project-1",
|
||||
project_metadata={"model_batch_enqueued_token_limit": {"routed-batch-model": 1}},
|
||||
)
|
||||
data, _ = await _model_routed_batch_data(openai_files, monkeypatch, query_string, routing_headers, user_api_key_dict)
|
||||
_, _, batch_limiter = _make_limiters()
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await batch_limiter.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type="acreate_batch"
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 429
|
||||
assert "project_model: project-1:routed-batch-model" in str(exc.value.detail)
|
||||
|
||||
|
||||
async def _install_managed_file_hook(
|
||||
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rows: str
|
||||
openai_files: respx.MockRouter,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
rows: Sequence[_BatchLine],
|
||||
file_id: str = _MANAGED_FILE_ID,
|
||||
model_name: str = _BATCH_MODEL,
|
||||
) -> tuple[respx.Route, respx.Route]:
|
||||
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
|
||||
|
||||
|
|
@ -642,18 +688,18 @@ async def _install_managed_file_hook(
|
|||
from litellm.proxy.openai_files_endpoints.common_utils import is_base64_encoded_unified_file_id
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
assert is_base64_encoded_unified_file_id(_MANAGED_FILE_ID)
|
||||
assert is_base64_encoded_unified_file_id(file_id)
|
||||
provider_route: Final = _serve_file(openai_files, "file-provider-1", rows)
|
||||
unrouted_managed_read: Final = _serve_file(openai_files, _MANAGED_FILE_ID, rows)
|
||||
unrouted_managed_read: Final = _serve_file(openai_files, file_id, rows)
|
||||
row: Final = LiteLLM_ManagedFileTable(
|
||||
unified_file_id=_MANAGED_FILE_ID,
|
||||
unified_file_id=file_id,
|
||||
model_mappings={"deployment-1": "file-provider-1"},
|
||||
flat_model_file_ids=["file-provider-1"],
|
||||
created_by="test-user-regression",
|
||||
team_id="team-owner",
|
||||
)
|
||||
file_cache: Final = InternalUsageCache(dual_cache=DualCache())
|
||||
await file_cache.async_set_cache(key=_MANAGED_FILE_ID, value=row.model_dump(), litellm_parent_otel_span=None)
|
||||
await file_cache.async_set_cache(key=file_id, value=row.model_dump(), litellm_parent_otel_span=None)
|
||||
prisma_client: Final = MagicMock()
|
||||
prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=row)
|
||||
proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache())
|
||||
|
|
@ -663,7 +709,7 @@ async def _install_managed_file_hook(
|
|||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": _BATCH_MODEL,
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"model": f"openai/{_BATCH_MODEL}", "api_key": "sk-test"},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
}
|
||||
|
|
@ -674,6 +720,72 @@ async def _install_managed_file_hook(
|
|||
return provider_route, unrouted_managed_read
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_managed_batch_file_charges_only_the_project_cap_of_the_model_it_routes_to(
|
||||
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import get_request_stash
|
||||
|
||||
priority_group: Final = f"{_BATCH_MODEL}-priority"
|
||||
managed_file_id: Final = base64.urlsafe_b64encode(
|
||||
f"litellm_proxy:application/octet-stream;unified_id,priority-batch;target_model_names,{priority_group}".encode()
|
||||
).decode()
|
||||
await _install_managed_file_hook(
|
||||
openai_files, monkeypatch, _batch_rows(("Hello",) * 3), file_id=managed_file_id, model_name=priority_group
|
||||
)
|
||||
_, _, batch_limiter = _make_limiters()
|
||||
data: Final = {"input_file_id": managed_file_id}
|
||||
|
||||
result: Final = await batch_limiter.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="test-key-regression",
|
||||
user_id="test-user-regression",
|
||||
team_id="team-owner",
|
||||
project_id="priority-project",
|
||||
project_metadata={"model_batch_enqueued_token_limit": {priority_group: 1000000, _BATCH_MODEL: 1}},
|
||||
),
|
||||
cache=DualCache(),
|
||||
data=data,
|
||||
call_type="acreate_batch",
|
||||
)
|
||||
|
||||
assert result is data
|
||||
stash: Final = get_request_stash()
|
||||
assert stash is not None and stash.batch_enqueued_reservation is not None
|
||||
assert [scope.value for scope in stash.batch_enqueued_reservation.scopes] == [f"priority-project:{priority_group}"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_project_cap_on_the_routed_model_lifts_per_minute_limits_off_a_managed_batch_file(
|
||||
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
priority_group: Final = f"{_BATCH_MODEL}-priority"
|
||||
managed_file_id: Final = base64.urlsafe_b64encode(
|
||||
f"litellm_proxy:application/octet-stream;unified_id,priority-minute;target_model_names,{priority_group}".encode()
|
||||
).decode()
|
||||
await _install_managed_file_hook(
|
||||
openai_files, monkeypatch, _batch_rows(("Hello",) * 3), file_id=managed_file_id, model_name=priority_group
|
||||
)
|
||||
_, _, batch_limiter = _make_limiters()
|
||||
data: Final = {"input_file_id": managed_file_id}
|
||||
|
||||
result: Final = await batch_limiter.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="test-key-regression",
|
||||
user_id="test-user-regression",
|
||||
team_id="team-owner",
|
||||
tpm_limit=1,
|
||||
project_id="priority-project",
|
||||
project_metadata={"model_batch_enqueued_token_limit": {priority_group: 1000000}},
|
||||
),
|
||||
cache=DualCache(),
|
||||
data=data,
|
||||
call_type="acreate_batch",
|
||||
)
|
||||
|
||||
assert result is data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_managed_file_is_read_through_the_managed_files_hook_with_user_context(
|
||||
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue