scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

feat(dynamic_rate_limiter_v3): priority_class_share_splits_equally_among_active_teams splits a priority class share among the teams active on a model (#51366)

Browse source 
* feat(dynamic_rate_limiter_v3): split a priority class share equally among the teams active on a model

Adds priority_reservation_settings.priority_class_share_splits_equally_among_active_teams (default false).
When on, the class weights are renormalized over the classes with a team active on the model in the
current window, each class share is divided equally among its active teams, and once the model saturates
each team is held to its own share with a 429 whose rate_limited_by header reads
priority_class_share_saturated. Active teams are one Redis sorted set per model with the window as TTL.

* fix(proxy): keep basedpyright counts flat for the priority class share setting

PriorityReservationSettings is built with model_validate instead of **value, so a new
field no longer adds an unknown-argument diagnostic, and the limiter unit test helpers
type the handler by its implementation class.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(dynamic_rate_limiter_v3): keep team shares usable when Redis, mixed-priority keys, WebSocket turns or tiny shares get in the way

A registry failure falls back to the class pool instead of skipping the model limits, a team's
counter is keyed by model and team so keys of different priorities spend one share, Responses
WebSocket turns are held to the class share like the team model limit, a share is floored at one
request and one token, and the registry declares the rate limits key family.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(dynamic_rate_limiter_v3): catch only Redis and validation errors when the active team registry fails

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(dynamic_rate_limiter_v3): keys without a team keep their class present, team ids may contain the registry separator, breaker-open counts as a registry failure

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(dynamic_rate_limiter_v3): per-class pool markers that take no team slice, collision-free model and team counter keys, local registry pruned by last-seen time

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(dynamic_rate_limiter_v3): settle tokens only on the stash the call owns, keep a configured class named default apart from teams with no priority

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(dynamic_rate_limiter_v3): tag team registry members and count every unenforced request against the team share

A team id can no longer spell a pool marker, since real teams are stored under a team prefix. Below saturation the team share counter is bumped through the atomic increment at the model-wide limit, so the local over-limit short-circuit never skips a request and the counter is complete when the share grows.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(dynamic_rate_limiter_v3): list the caller under its own class and resync the socket stash share on every frame

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(dynamic_rate_limiter_v3): keep the connection share across a registry outage and settle fallback tokens on the served model's team counter

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-11-12 06:38:07 -08:00 • committed by GitHub
parent 52d1c4cd80
commit c0114e55aa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 1074 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

199
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -11,12 +11,14 @@ from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from fastapi import HTTPException
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from redis.exceptions import RedisError
import litellm
from litellm import Router
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
from litellm.cost_calculator import get_usage_object
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.add_retry_fallback_headers import (
@ -37,13 +39,25 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401 # l
RateLimitedModel,
RateLimitResponse,
RateLimitStatus,
_call_id_from_callback_kwargs, # pyright: ignore[reportPrivateUsage] # same call-ownership check the regular limiter uses
_is_socket_turn, # pyright: ignore[reportPrivateUsage] # same socket-turn check the regular limiter uses
_PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_socket_turn_tpm_windows, # pyright: ignore[reportPrivateUsage] # same per-turn TPM window the regular limiter checks
borrowing_team_model_limits,
claim_request_stash_for_data,
get_or_create_request_stash,
get_request_stash,
get_request_stash_for_call,
)
from litellm.proxy.hooks.priority_class_share import (
DEFAULT_CLASS,
ActiveTeamRegistry,
TeamShare,
counter_value,
pool_member,
team_member,
team_share,
)
from litellm.proxy.hooks.rate_limiter_utils import (
convert_priority_to_percent,
resolve_llm_provider_for_rate_limit,
@ -341,6 +355,7 @@ def _usage_of(response_obj: object) -> Usage | None:
TEAM_MODEL_LIMIT_AT_SATURATION: Final = "model_per_team_saturated"
PRIORITY_CLASS_SHARE_AT_SATURATION: Final = "priority_class_share_saturated"
TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE: Final = "model_per_team_max_borrow_multiple"
@ -427,6 +442,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
self.internal_usage_cache = InternalUsageCache(dual_cache=internal_usage_cache)
self._is_model_authorized = is_model_authorized
self.v3_limiter = PROXY_MaxParallelRequestsHandler_v3(self.internal_usage_cache, time_provider=time_provider)
self.active_teams = ActiveTeamRegistry(internal_usage_cache)
def update_variables(self, llm_router: Router) -> None:
self.llm_router = llm_router
@ -619,11 +635,54 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
# Fail open: assume not saturated on error
return 0.0
def _priority_class(self, priority: str | None) -> str:
"""The class a team is counted under: its explicit priority, or one shared default class."""
if (
priority is not None
and litellm.priority_reservation is not None
and priority in litellm.priority_reservation
):
return priority
return DEFAULT_CLASS
async def _team_class_share(
self,
model: str,
user_api_key_dict: UserAPIKeyAuth,
priority: str | None,
normalized_weights: Mapping[str, float],
) -> TeamShare | None:
"""The caller's slice of its class share, or None when the split is off or the key has no team. A key with no
team keeps the class pool but still marks its class active, so the pool's capacity is not lent away."""
if (
litellm.priority_reservation is None
or not _get_priority_settings().priority_class_share_splits_equally_among_active_teams
):
return None
priority_class: Final = self._priority_class(priority)
member: Final = (
pool_member(priority_class) if user_api_key_dict.team_id is None else team_member(user_api_key_dict.team_id)
)
try:
members: Final = await self.active_teams.touch(model, member, priority_class, self.v3_limiter.window_size)
except (RedisError, RedisCircuitBreakerOpenError, ValidationError) as e:
verbose_proxy_logger.warning(
"Active team registry unavailable for %s, holding team %s to the class pool: %s",
model,
user_api_key_dict.team_id,
e,
)
return None
if user_api_key_dict.team_id is None:
return None
return team_share(priority_class, members, normalized_weights, _get_priority_settings().default_priority)
def _create_priority_based_descriptors(
self,
model: str,
user_api_key_dict: UserAPIKeyAuth,
priority: str | None,
share: TeamShare | None = None,
) -> list[RateLimitDescriptor]:
"""
Create rate limit descriptors with normalized priority weights.
@ -653,14 +712,16 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
)
rate_limit_config: Final[RateLimitDescriptorRateLimitObject] = {}
weight: Final = priority_weight if share is None else share.fraction
floor: Final = 0 if share is None else 1
# Apply priority weight to model limits
if model_group_info.tpm is not None:
reserved_tpm: Final = int(model_group_info.tpm * priority_weight)
reserved_tpm: Final = max(floor, int(model_group_info.tpm * weight))
rate_limit_config["tokens_per_unit"] = reserved_tpm
if model_group_info.rpm is not None:
reserved_rpm: Final = int(model_group_info.rpm * priority_weight)
reserved_rpm: Final = max(floor, int(model_group_info.rpm * weight))
rate_limit_config["requests_per_unit"] = reserved_rpm
if rate_limit_config:
@ -668,8 +729,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
descriptors.append(
RateLimitDescriptor(
key="priority_model",
value=priority_key,
key="priority_model" if share is None else PRIORITY_CLASS_SHARE_AT_SATURATION,
value=priority_key if share is None else counter_value(model, str(user_api_key_dict.team_id)),
rate_limit=rate_limit_config,
)
)
@ -783,35 +844,93 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
llm_provider=llm_provider,
)
def _class_share_error(
self,
model: str,
status: RateLimitStatus,
saturation: float,
priority: str | None,
user_api_key_dict: UserAPIKeyAuth,
share: TeamShare,
) -> ProxyRateLimitError:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model)
return ProxyRateLimitError(
detail={
"error": f"Priority class share reached for {model} and the model is saturated. "
f"Team: {user_api_key_dict.team_id}, "
f"Priority: {priority or 'default'}, "
f"Active teams in class: {share.active_teams_in_class}, "
f"Active classes: {share.active_classes}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Team share: {status['current_limit']}, "
f"Remaining: {status['limit_remaining']}, "
f"Model saturation: {saturation:.1%}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
"rate_limit_type": str(status["rate_limit_type"]),
"rate_limited_by": PRIORITY_CLASS_SHARE_AT_SATURATION,
"x-litellm-priority": priority or "default",
"x-litellm-saturation": f"{saturation:.2%}",
},
rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]),
model=resolved_model,
llm_provider=llm_provider,
)
async def enforce_token_limits_for_frame(
self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None
) -> None:
"""Hold a Responses WebSocket turn to the team's share once the model saturates, like the regular limiter does."""
"""Hold a Responses WebSocket turn to the team's shares once the model saturates, like the regular limiter does."""
if requested_model is None:
return
team_share: Final = self._borrowing_team_share(model=requested_model, user_api_key_dict=user_api_key_dict)
model_group_info: Final = self.llm_router.registry().group(requested_model)
if team_share is None or model_group_info is None:
if model_group_info is None:
return
priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict)
share: Final = await self._team_class_share(
requested_model, user_api_key_dict, priority, self._normalize_priority_weights(model_group_info)
)
share_descriptors: Final = (
()
if share is None
else tuple(
self._create_priority_based_descriptors(
model=requested_model, user_api_key_dict=user_api_key_dict, priority=priority, share=share
)
)
)
frame_share: Final = next(
(d for d in share_descriptors if d["key"] == PRIORITY_CLASS_SHARE_AT_SATURATION), None
)
connection_stash: Final = get_request_stash()
if connection_stash is not None and frame_share is not None:
connection_stash.priority_class_share = frame_share
borrowing: Final = self._borrowing_team_share(model=requested_model, user_api_key_dict=user_api_key_dict)
if not share_descriptors and borrowing is None:
return
saturation: Final = await self._check_model_saturation(requested_model, model_group_info)
enforced: Final = team_share.enforced(saturation >= _get_priority_settings().saturation_threshold)
if enforced is None:
saturated: Final = saturation >= _get_priority_settings().saturation_threshold
enforced_team: Final = None if borrowing is None else borrowing.enforced(saturated)
descriptors: Final = (share_descriptors if saturated else ()) + (
() if enforced_team is None else (enforced_team,)
)
if not descriptors:
return
tpm_windows: Final = _socket_turn_tpm_windows((enforced,))
tpm_windows: Final = _socket_turn_tpm_windows(descriptors)
response: Final = await self.v3_limiter.should_rate_limit(
tpm_windows, parent_otel_span=user_api_key_dict.parent_otel_span, read_only=True, skip_parallel_check=True
)
over_limit: Final = next((s for s in response["statuses"] if s["code"] == "OVER_LIMIT"), None)
if over_limit is None:
return
raise self._team_share_error(
requested_model,
over_limit,
saturation,
self._get_priority_from_user_api_key_dict(user_api_key_dict),
user_api_key_dict,
team_share,
)
over_class_share: Final = over_limit["descriptor_key"] == PRIORITY_CLASS_SHARE_AT_SATURATION
if share is not None and (over_class_share or borrowing is None):
raise self._class_share_error(requested_model, over_limit, saturation, priority, user_api_key_dict, share)
if borrowing is not None:
raise self._team_share_error(
requested_model, over_limit, saturation, priority, user_api_key_dict, borrowing
)
def _create_model_tracking_descriptor(
self,
@ -893,10 +1012,17 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
descriptors_to_check.append(model_wide_descriptor)
# Priority descriptors (always track, conditionally enforce)
share: Final = await self._team_class_share(
model, user_api_key_dict, priority, self._normalize_priority_weights(model_group_info)
)
priority_descriptors: Final = self._create_priority_based_descriptors(
model=model,
user_api_key_dict=user_api_key_dict,
priority=priority,
share=share,
)
get_or_create_request_stash().priority_class_share = next(
(d for d in priority_descriptors if d["key"] == PRIORITY_CLASS_SHARE_AT_SATURATION), None
)
if priority_descriptors:
descriptors_to_check.extend(priority_descriptors)
@ -986,6 +1112,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
if descriptor_key == TEAM_MODEL_LIMIT_AT_SATURATION and team_share is not None:
raise self._team_share_error(model, status, saturation, priority, user_api_key_dict, team_share)
if descriptor_key == PRIORITY_CLASS_SHARE_AT_SATURATION and share is not None:
raise self._class_share_error(model, status, saturation, priority, user_api_key_dict, share)
# Fail-closed guard: overall_code says OVER_LIMIT but no status
# matched a descriptor key we know how to translate into a 429.
@ -1032,7 +1160,22 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
increments=[per_request_increment],
parent_otel_span=user_api_key_dict.parent_otel_span,
)
if priority_descriptors and not should_enforce_priority:
if priority_descriptors and share is not None and not should_enforce_priority:
tracking_only: Final = tuple(
RateLimitDescriptor(
key=descriptor["key"],
value=descriptor["value"],
rate_limit=model_wide_descriptor["rate_limit"],
)
for descriptor in priority_descriptors
)
await self.v3_limiter.atomic_check_and_increment_by_n(
descriptors=list(tracking_only),
increments=[per_request_increment for _ in tracking_only],
parent_otel_span=user_api_key_dict.parent_otel_span,
)
get_or_create_request_stash().rate_limit_response = atomic_response
elif priority_descriptors and not should_enforce_priority:
priority_tracking_response: Final = await self.v3_limiter.should_rate_limit(
descriptors=priority_descriptors,
parent_otel_span=user_api_key_dict.parent_otel_span,
@ -1378,13 +1521,27 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
else:
priority_key = f"{model_group}:default_pool"
call_stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs))
stash: Final = (
call_stash if call_stash is not None else (get_request_stash() if _is_socket_turn(kwargs) else None)
)
team_model_tokens: Final = _team_model_token_increment(
get_or_create_request_stash().borrowing_team_model_limit, total_tokens, self.v3_limiter.window_size
None if stash is None else stash.borrowing_team_model_limit, total_tokens, self.v3_limiter.window_size
)
class_share: Final = None if stash is None else stash.priority_class_share
settled_team_id: Final = settlement_metadata.user_api_key_team_id
priority_scope: Final = (
RateLimitScope("priority_model", priority_key)
if class_share is None
else RateLimitScope(
class_share["key"],
class_share["value"] if settled_team_id is None else counter_value(model_group, settled_team_id),
)
)
pipeline_operations: Final = (
model_tokens,
CounterIncrement(
counter=RateLimitScope("priority_model", priority_key).counter("tokens"),
counter=priority_scope.counter("tokens"),
amount=total_tokens,
ttl=self.v3_limiter.window_size,
),

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

@ -615,6 +615,7 @@ class RequestRateLimiterStash:
batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None
batch_tpd_refund_ops: tuple[ReservationOperation, ...] = ()
borrowing_team_model_limit: RateLimitDescriptor | None = None
priority_class_share: RateLimitDescriptor | None = None
reservation_released: bool = False
io_reservation_released: bool = False
tpm_limited_tags: frozenset[str] = field(default_factory=frozenset)

127
litellm/proxy/hooks/priority_class_share.py Normal file
View file

@ -0,0 +1,127 @@
"""Splits a priority class's share of a model among the teams active on it.
Enabled by ``priority_reservation_settings.priority_class_share_splits_equally_among_active_teams``.
"""
import time
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from pydantic import TypeAdapter
from litellm._internal_context import with_service_target
from litellm.caching.caching import DualCache
from litellm.redis.stores.rate_limits import RATE_LIMITS_TARGET
ACTIVE_TEAMS_KEY_PREFIX: Final = "dynamic_rate_limiter:active_teams:"
MEMBER_SEPARATOR: Final = "\x1f"
DEFAULT_CLASS: Final = "\x1edefault"
_POOL_PREFIX: Final = f"{MEMBER_SEPARATOR}pool{MEMBER_SEPARATOR}"
_TOUCH_LUA: Final = (
"redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) "
"redis.call('ZADD', KEYS[1], ARGV[2], ARGV[3]) "
"redis.call('EXPIRE', KEYS[1], ARGV[4]) "
"return redis.call('ZRANGE', KEYS[1], 0, -1)"
)
_MEMBERS: Final = TypeAdapter[tuple[str | bytes, ...]](tuple[str | bytes, ...])
def pool_member(priority: str) -> str:
"""The registry id that stands for every key without a team in ``priority``. It keeps the class present without
being a team, and starts with the separator so no team id can spell it."""
return f"{_POOL_PREFIX}{priority}"
def team_member(team_id: str) -> str:
"""The registry id of a real team. It never starts with the separator, so it cannot spell a pool marker."""
return f"team{MEMBER_SEPARATOR}{team_id}"
def is_pool_member(member: str) -> bool:
return member.startswith(_POOL_PREFIX)
def counter_value(model: str, team_id: str) -> str:
"""One counter per model and team. The length prefix keeps ``chat`` + ``prod:victim`` apart from
``chat:prod`` + ``victim``."""
return f"{len(model)}:{model}:{team_id}"
@dataclass(frozen=True, slots=True)
class TeamShare:
fraction: float
active_teams_in_class: int
active_classes: int
def team_share(
priority: str, members: Mapping[str, str], class_weights: Mapping[str, float], default_weight: float
) -> TeamShare:
"""``members`` maps every team active on the model, the caller included, to its priority class. Classes with no
active member give their weight away, and each present class is split equally among its active teams. A pool
member keeps its class present but takes no slice, since keys without a team are held to the class pool."""
present: Final = frozenset(members.values())
weights: Final = {cls: class_weights.get(cls, default_weight) for cls in present}
total: Final = sum(weights.values())
class_fraction: Final = weights[priority] / total if total > 0 else 1.0 / len(present)
teams_in_class: Final = sum(1 for team, cls in members.items() if cls == priority and not is_pool_member(team))
return TeamShare(
fraction=class_fraction / max(1, teams_in_class),
active_teams_in_class=teams_in_class,
active_classes=len(present),
)
def _members(raw: tuple[str | bytes, ...]) -> Mapping[str, str]:
"""Oldest first, so a team seen under two classes in one window keeps its most recent one. The class is split off
from the right because a team id may itself contain the separator."""
decoded: Final = (item.decode() if isinstance(item, bytes) else item for item in raw)
pairs: Final = (member.rsplit(MEMBER_SEPARATOR, 1) for member in decoded)
return MappingProxyType({pair[0]: pair[1] for pair in pairs if len(pair) == 2})
class ActiveTeamRegistry:
"""Which teams used a model in the current window, as one sorted set per model scored by last-seen time, in Redis
when the proxy has one and otherwise in process memory."""
def __init__(self, dual_cache: DualCache, time_provider: Callable[[], float] = time.time) -> None:
self._dual_cache = dual_cache
self._local: Mapping[str, Mapping[str, float]] = MappingProxyType({})
self._now = time_provider
@with_service_target(RATE_LIMITS_TARGET)
async def touch(self, model: str, team_id: str, priority: str, window: int) -> Mapping[str, str]:
"""Record ``team_id`` as active on ``model`` and return every team active on it, the caller included. The
caller is always listed under ``priority``, whatever class another worker's clock last saw it in."""
key: Final = f"{ACTIVE_TEAMS_KEY_PREFIX}{model}"
member: Final = f"{team_id}{MEMBER_SEPARATOR}{priority}"
now: Final = self._now()
redis_cache: Final = self._dual_cache.redis_cache
seen: Final = (
self._touch_in_memory(key, member, now, window)
if redis_cache is None
else _members(
_MEMBERS.validate_python(
await redis_cache.async_register_script(_TOUCH_LUA)(
keys=(key,), args=(now - window, now, member, window)
)
)
)
)
return MappingProxyType({**seen, team_id: priority})
def _touch_in_memory(self, key: str, member: str, now: float, window: int) -> Mapping[str, str]:
"""Prunes by each member's own last-seen time, so the map lives exactly as long as its newest member."""
cutoff: Final = now - window
alive: Final = {seen_member: seen for seen_member, seen in self._local.get(key, {}).items() if seen > cutoff}
current: Final = MappingProxyType({**alive, member: now})
others: Final = {
other_key: others_members
for other_key, others_members in self._local.items()
if other_key != key and any(seen > cutoff for seen in others_members.values())
}
self._local = MappingProxyType({**others, key: current})
return _members(tuple(sorted(current, key=current.__getitem__)))

2
litellm/proxy/proxy_server.py
View file

@ -7080,7 +7080,7 @@ class ProxyConfig:
elif key == "priority_reservation_settings":
from litellm.types.utils import PriorityReservationSettings
litellm.priority_reservation_settings = PriorityReservationSettings(**value)
litellm.priority_reservation_settings = PriorityReservationSettings.model_validate(value)
elif key == "callbacks":
initialize_callbacks_on_proxy(
value=value,

4
litellm/types/utils.py
View file

@ -5393,6 +5393,10 @@ class PriorityReservationSettings(LiteLLMBaseModel):
default=False,
description="When true, a team's model_tpm_limit and model_rpm_limit on a model this limiter covers stop being ceilings. The team may run past them on capacity other teams are not using while the model is below saturation_threshold, and is held to them once the model saturates.",
)
priority_class_share_splits_equally_among_active_teams: bool = Field(
default=False,
description="When true, a priority class's share of a model is renormalized over the classes whose teams used the model in the current window and split equally among that class's active teams, and once the model saturates each team is held to its own share instead of the class pool. Keys without a team keep the class pool.",
)
saturation_check_cache_ttl: int = Field(
default=60,

136
tests/integration/routing/test_priority_class_share_splits_among_active_teams.py Normal file
View file

@ -0,0 +1,136 @@
import json
import uuid
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, Scenario, gateway_from_environment
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
OPENAI_MODEL: Final = "gpt-4.1-mini"
MODEL_TPM: Final = 1000
TOKENS_PER_CALL: Final = 30
UPSTREAM_REPLY: Final = json.dumps(
{
"id": "chatcmpl_priority_class_share",
"object": "chat.completion",
"created": 1700000000,
"model": OPENAI_MODEL,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "shared"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": TOKENS_PER_CALL},
}
).encode()
def _respond(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/v1/chat/completions", request.target
return Reply(body=UPSTREAM_REPLY)
@pytest.fixture(scope="module")
def saturated_model(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, str]]:
directory: Final = tmp_path_factory.mktemp("priority-class-share")
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["litellm_settings"] = {
**configuration["litellm_settings"],
"callbacks": ["dynamic_rate_limiter_v3"],
"priority_reservation": {"interactive": 0.6, "batch": 0.4},
"priority_reservation_settings": {
"default_priority": 1.0,
"saturation_threshold": 0.0,
"priority_class_share_splits_equally_among_active_teams": True,
},
}
path: Final = directory / "priority_class_share.yaml"
path.write_text(yaml.safe_dump(configuration))
with (
gateway_from_environment() as rig,
wire_server(_respond) as wire,
owned_proxy(rig, directory, {}, config=path) as owned,
):
yield owned, wire.url
def _team_key(scenario: Scenario, priority: str) -> str:
return scenario.key(team_id=scenario.team(metadata={"priority": priority}))
def _call(owned: Gateway, model: str, key: str) -> httpx.Response:
return owned.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"class share {uuid.uuid4().hex}"}]},
key=key,
)
def _spend_until_blocked(owned: Gateway, model: str, key: str, max_calls: int) -> tuple[httpx.Response, ...]:
def calls() -> Iterator[httpx.Response]:
for _ in range(max_calls):
response: Final = _call(owned, model, key)
yield response
if response.status_code == 429:
return
return tuple(calls())
def test_second_interactive_team_keeps_its_half_of_the_class_share(saturated_model: tuple[Gateway, str]) -> None:
"""Two interactive teams, batch idle: the class share is the whole model and each team is held to half of it."""
owned, upstream = saturated_model
with owned.scenario() as scenario:
model: Final = scenario.model(
model=f"openai/{OPENAI_MODEL}", api_base=f"{upstream}/v1", api_key="synthetic-openai-key", tpm=MODEL_TPM
)
first: Final = _team_key(scenario, "interactive")
second: Final = _team_key(scenario, "interactive")
assert _call(owned, model, second).status_code == 200
responses: Final = _spend_until_blocked(owned, model, first, max_calls=MODEL_TPM // TOKENS_PER_CALL + 1)
blocked: Final = responses[-1]
assert blocked.status_code == 429, blocked.text
assert blocked.headers.get("rate_limited_by") == "priority_class_share_saturated", dict(blocked.headers)
served: Final = (len(responses) - 1) * TOKENS_PER_CALL
assert MODEL_TPM // 2 <= served < MODEL_TPM // 2 + 2 * TOKENS_PER_CALL, served
still_served: Final = _call(owned, model, second)
assert still_served.status_code == 200, still_served.text
def test_only_active_class_gets_the_whole_model(saturated_model: tuple[Gateway, str]) -> None:
"""One batch team alone on the model is held to 100% of the model, not to batch's 40%."""
owned, upstream = saturated_model
with owned.scenario() as scenario:
model: Final = scenario.model(
model=f"openai/{OPENAI_MODEL}", api_base=f"{upstream}/v1", api_key="synthetic-openai-key", tpm=MODEL_TPM
)
batch: Final = _team_key(scenario, "batch")
responses: Final = _spend_until_blocked(owned, model, batch, max_calls=MODEL_TPM // TOKENS_PER_CALL + 1)
served: Final = sum(1 for response in responses if response.status_code == 200) * TOKENS_PER_CALL
assert served >= MODEL_TPM - 2 * TOKENS_PER_CALL, [r.text for r in responses]
def test_class_weights_renormalize_when_batch_wakes_up(saturated_model: tuple[Gateway, str]) -> None:
"""Interactive and batch both active: interactive is held to 60% of the model and batch to 40%."""
owned, upstream = saturated_model
with owned.scenario() as scenario:
model: Final = scenario.model(
model=f"openai/{OPENAI_MODEL}", api_base=f"{upstream}/v1", api_key="synthetic-openai-key", tpm=MODEL_TPM
)
interactive: Final = _team_key(scenario, "interactive")
batch: Final = _team_key(scenario, "batch")
assert _call(owned, model, batch).status_code == 200
responses: Final = _spend_until_blocked(owned, model, interactive, max_calls=MODEL_TPM // TOKENS_PER_CALL + 1)
assert responses[-1].status_code == 429, responses[-1].text
served: Final = (len(responses) - 1) * TOKENS_PER_CALL
assert int(MODEL_TPM * 0.6) <= served < int(MODEL_TPM * 0.6) + 2 * TOKENS_PER_CALL, served
batch_responses: Final = _spend_until_blocked(owned, model, batch, max_calls=MODEL_TPM // TOKENS_PER_CALL + 1)
batch_served: Final = sum(1 for response in batch_responses if response.status_code == 200) * TOKENS_PER_CALL
assert batch_served >= int(MODEL_TPM * 0.4) - 2 * TOKENS_PER_CALL, [r.text for r in batch_responses]
assert batch_responses[-1].status_code == 429, batch_responses[-1].text

478
tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py
View file

@ -22,6 +22,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
AlternateGrant,
PROXY_DynamicRateLimitHandlerV3,
_PROXY_DynamicRateLimitHandlerV3, # pyright: ignore[reportPrivateUsage] # the public name is a Final alias, not a type
)
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
PROXY_DynamicRateLimitHandlerV3 as DynamicRateLimitHandler,
@ -29,6 +30,7 @@ from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.hooks.priority_class_share import pool_member
from litellm.redis.stores.rate_limits import RateLimitScope
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits
@ -2114,7 +2116,9 @@ def _borrowing_handler(
borrows: bool,
saturation_threshold: float = 0.5,
model_group_alias: dict[str, str] | None = None,
) -> PROXY_DynamicRateLimitHandlerV3:
splits: bool = False,
rpm: int = 100,
) -> _PROXY_DynamicRateLimitHandlerV3:
from litellm.types.utils import PriorityReservationSettings
monkeypatch.setenv("LITELLM_LICENSE", "test-license-key")
@ -2123,7 +2127,9 @@ def _borrowing_handler(
litellm,
"priority_reservation_settings",
PriorityReservationSettings(
saturation_threshold=saturation_threshold, team_model_limit_borrows_idle_capacity_from_other_teams=borrows
saturation_threshold=saturation_threshold,
team_model_limit_borrows_idle_capacity_from_other_teams=borrows,
priority_class_share_splits_equally_among_active_teams=splits,
),
)
monkeypatch.setattr(litellm, "callbacks", ["dynamic_rate_limiter_v3"])
@ -2131,8 +2137,20 @@ def _borrowing_handler(
model_list=[
{
"model_name": "shared-model",
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "k", "api_base": "b", "tpm": 1000, "rpm": rpm},
},
{
"model_name": "other-model",
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "k", "api_base": "b", "tpm": 1000, "rpm": 100},
}
},
{
"model_name": "chat",
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "k", "api_base": "b", "tpm": 1000, "rpm": 100},
},
{
"model_name": "chat:prod",
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "k", "api_base": "b", "tpm": 1000, "rpm": 100},
},
],
model_group_alias=model_group_alias or {},
)
@ -2142,12 +2160,14 @@ def _borrowing_handler(
return handler
async def _admit_then_settle(handler: PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, tokens: int) -> None:
async def _admit_then_settle(handler: _PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, tokens: int) -> None:
await _check(handler, auth, saturation=0.0)
await _settle_tokens(handler, tokens)
async def _settle_tokens(handler: PROXY_DynamicRateLimitHandlerV3, tokens: int, response: object | None = None) -> None:
async def _settle_tokens(
handler: _PROXY_DynamicRateLimitHandlerV3, tokens: int, response: object | None = None
) -> None:
await handler.async_log_success_event(
kwargs={
"litellm_params": {"metadata": {"model_group": "shared-model"}},
@ -2177,12 +2197,16 @@ def _team_key(
async def _check(
handler: PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, saturation: float, model: str = "shared-model"
handler: _PROXY_DynamicRateLimitHandlerV3,
auth: UserAPIKeyAuth,
saturation: float,
model: str = "shared-model",
priority: str | None = None,
) -> None:
group: Final = handler.llm_router.registry().group(model)
assert group is not None
await handler._check_rate_limits(
model=model, model_group_info=group, user_api_key_dict=auth, priority=None, saturation=saturation
await handler._check_rate_limits( # pyright: ignore[reportPrivateUsage] # saturation is an input here, the public hook reads it from Redis
model=model, model_group_info=group, user_api_key_dict=auth, priority=priority, saturation=saturation
)
@ -2367,6 +2391,273 @@ async def test_a_websocket_turn_still_borrows_while_the_model_is_idle(monkeypatc
await handler.enforce_token_limits_for_frame(user_api_key_dict=_team_key(), requested_model="shared-model")
def _class_team(team_id: str, priority: str | None = "interactive") -> UserAPIKeyAuth:
return UserAPIKeyAuth(api_key="sk-class", team_id=team_id, team_metadata={}, metadata={"priority": priority})
async def _burn(
handler: _PROXY_DynamicRateLimitHandlerV3,
auth: UserAPIKeyAuth,
priority: str | None,
tokens: int,
model: str = "shared-model",
):
await _check(handler, auth, saturation=0.0, priority=priority, model=model)
await handler.async_log_success_event(
kwargs={
"litellm_params": {"metadata": {"model_group": model}},
"standard_logging_object": {
"metadata": {
"user_api_key_team_id": auth.team_id,
"user_api_key_auth_metadata": {"priority": priority},
}
},
},
response_obj=litellm.ModelResponse(
usage=litellm.Usage(prompt_tokens=tokens, completion_tokens=0, total_tokens=tokens)
),
start_time=datetime.now(),
end_time=datetime.now(),
)
def _held_to_class_share(exc: ProxyRateLimitError) -> None:
assert exc.status_code == 429
assert exc.headers is not None
assert exc.headers["rate_limited_by"] == "priority_class_share_saturated"
assert exc.headers["rate_limit_type"] == "tokens"
@pytest.mark.asyncio
async def test_two_interactive_teams_each_keep_half_of_the_interactive_share(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team("team-a"), "interactive", 20)
await _burn(handler, _class_team("team-b"), "interactive", 300)
await _burn(handler, _class_team("team-b"), "interactive", 200)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-b"), saturation=0.52, priority="interactive")
_held_to_class_share(exc.value)
assert "Team: team-b, Priority: interactive, Active teams in class: 2, Active classes: 1, " in str(exc.value.detail)
assert "Team share: 500, Remaining: 0, Model saturation: 52.0%" in str(exc.value.detail)
await _check(handler, _class_team("team-a"), saturation=0.52, priority="interactive")
@pytest.mark.asyncio
async def test_the_only_active_class_gets_the_whole_model_instead_of_its_weight(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team("team-b", "batch"), "batch", 950)
await _check(handler, _class_team("team-b", "batch"), saturation=0.95, priority="batch")
unsplit = _borrowing_handler(monkeypatch, borrows=False)
await _burn(unsplit, _class_team("team-b", "batch"), "batch", 950)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(unsplit, _class_team("team-b", "batch"), saturation=0.95, priority="batch")
assert "Priority-based rate limit exceeded" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_class_weights_apply_again_once_a_team_of_the_other_class_is_active(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team("team-b", "batch"), "batch", 10)
await _burn(handler, _class_team("team-a"), "interactive", 600)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-a"), saturation=0.61, priority="interactive")
assert "Active teams in class: 1, Active classes: 2, Rate limit type: tokens, Team share: 600, " in str(
exc.value.detail
)
await _check(handler, _class_team("team-b", "batch"), saturation=0.61, priority="batch")
@pytest.mark.asyncio
async def test_teams_without_a_priority_split_the_default_share_between_them(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team("team-x", None), None, 600)
await _burn(handler, _class_team("team-y", None), None, 1)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-x", None), saturation=0.6, priority=None)
assert "Priority: default, Active teams in class: 2, Active classes: 1, " in str(exc.value.detail)
assert "Team share: 500, " in str(exc.value.detail)
await _check(handler, _class_team("team-y", None), saturation=0.6, priority=None)
assert exc.value.headers is not None and exc.value.headers["x-litellm-priority"] == "default"
@pytest.mark.asyncio
async def test_a_key_without_a_team_keeps_the_shared_class_pool(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team("team-a"), "interactive", 100)
no_team: Final = UserAPIKeyAuth(api_key="sk-plain", metadata={"priority": "interactive"})
await _burn(handler, no_team, "interactive", 600)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, no_team, saturation=0.7, priority="interactive")
assert exc.value.headers is not None and "rate_limited_by" not in exc.value.headers
assert "Priority-based rate limit exceeded" in str(exc.value.detail)
await _check(handler, _class_team("team-a"), saturation=0.7, priority="interactive")
@pytest.mark.asyncio
async def test_below_saturation_a_team_may_run_past_its_share(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team("team-a"), "interactive", 1)
await _burn(handler, _class_team("team-b"), "interactive", 400)
await _check(handler, _class_team("team-b"), saturation=0.49, priority="interactive")
@pytest.mark.asyncio
async def test_the_split_is_off_by_default_so_one_team_can_still_empty_the_class_pool(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False)
await _burn(handler, _class_team("team-a"), "interactive", 600)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-b"), saturation=0.6, priority="interactive")
assert exc.value.headers is not None and "rate_limited_by" not in exc.value.headers
assert "Priority-based rate limit exceeded" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_the_team_model_limit_and_the_class_share_both_apply_when_both_are_on(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=True, splits=True)
team: Final = UserAPIKeyAuth(
api_key="sk-both",
team_id="team-a",
team_metadata={"model_tpm_limit": {"shared-model": 100}},
metadata={"priority": "interactive"},
)
await _burn(handler, team, "interactive", 150)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, team, saturation=0.5, priority="interactive")
_held_to_team_limit(exc.value, "tokens", "model_tpm_limit")
@pytest.mark.asyncio
async def test_a_team_active_on_another_model_does_not_shrink_this_models_shares(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _check(handler, _class_team("team-b"), saturation=0.0, priority="interactive", model="other-model")
await _burn(handler, _class_team("team-a"), "interactive", 700)
await _check(handler, _class_team("team-a"), saturation=0.7, priority="interactive")
class _UnavailableRegistry:
async def touch(self, model: str, team_id: str, priority: str, window: int) -> Mapping[str, str]:
from redis.exceptions import ConnectionError as RedisConnectionError
raise RedisConnectionError("redis is down")
@pytest.mark.asyncio
async def test_a_registry_failure_falls_back_to_the_class_pool_instead_of_skipping_the_limits(
monkeypatch: pytest.MonkeyPatch,
):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
handler.active_teams = _UnavailableRegistry() # pyright: ignore[reportAttributeAccessIssue] # injected failing registry
await _burn(handler, _class_team("team-a"), "interactive", 650)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-a"), saturation=0.65, priority="interactive")
assert exc.value.headers is not None
assert "rate_limited_by" not in exc.value.headers
assert "Priority-based rate limit exceeded" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_one_team_using_keys_of_two_priorities_still_spends_one_share(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team("team-b"), "interactive", 20)
await _burn(handler, _class_team("team-a", "interactive"), "interactive", 300)
await _burn(handler, _class_team("team-a", "batch"), "batch", 300)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-a", "batch"), saturation=0.62, priority="batch")
_held_to_class_share(exc.value)
assert "Team: team-a, Priority: batch, Active teams in class: 1, Active classes: 2, " in str(exc.value.detail)
assert "Team share: 400, Remaining: 0," in str(exc.value.detail)
@pytest.mark.asyncio
async def test_a_share_smaller_than_one_request_still_lets_each_team_send_one(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True, rpm=2)
for team in ("team-a", "team-b", "team-c"):
await handler.active_teams.touch("shared-model", team, "interactive", 60)
await _check(handler, _class_team("team-d"), saturation=0.9, priority="interactive")
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-d"), saturation=0.9, priority="interactive")
assert exc.value.headers is not None
assert exc.value.headers["rate_limited_by"] == "priority_class_share_saturated"
assert exc.value.headers["rate_limit_type"] == "requests"
@pytest.mark.asyncio
async def test_a_websocket_turn_is_held_to_the_class_share_once_the_model_saturates(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True, saturation_threshold=0.1)
await _burn(handler, _class_team("team-a"), "interactive", 20)
await _burn(handler, _class_team("team-b"), "interactive", 600)
with pytest.raises(ProxyRateLimitError) as exc:
await handler.enforce_token_limits_for_frame(
user_api_key_dict=_class_team("team-b"), requested_model="shared-model"
)
_held_to_class_share(exc.value)
await handler.enforce_token_limits_for_frame(
user_api_key_dict=_class_team("team-a"), requested_model="shared-model"
)
@pytest.mark.asyncio
async def test_a_websocket_turn_ignores_the_class_share_while_the_split_is_off(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, saturation_threshold=0.1)
await _burn(handler, _class_team("team-a"), "interactive", 20)
await _burn(handler, _class_team("team-b"), "interactive", 600)
await handler.enforce_token_limits_for_frame(
user_api_key_dict=_class_team("team-b"), requested_model="shared-model"
)
@pytest.mark.asyncio
async def test_traffic_from_keys_without_a_team_keeps_its_class_present_in_the_split(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
no_team: Final = UserAPIKeyAuth(api_key="sk-plain", metadata={"priority": "interactive"})
await _burn(handler, no_team, "interactive", 20)
await _burn(handler, _class_team("team-c", "batch"), "batch", 450)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-c", "batch"), saturation=0.6, priority="batch")
_held_to_class_share(exc.value)
assert "Team share: 400, Remaining: 0," in str(exc.value.detail)
@pytest.mark.asyncio
async def test_a_team_id_containing_the_registry_separator_is_still_held_to_its_share(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team("team-b"), "interactive", 20)
await _burn(handler, _class_team("owned\x1fextra", "batch"), "batch", 450)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("owned\x1fextra", "batch"), saturation=0.6, priority="batch")
_held_to_class_share(exc.value)
@pytest.mark.asyncio
async def test_a_team_share_is_not_shrunk_by_teamless_traffic_in_its_class(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
no_team: Final = UserAPIKeyAuth(api_key="sk-plain", metadata={"priority": "interactive"})
await _burn(handler, no_team, "interactive", 20)
await _burn(handler, _class_team("team-a"), "interactive", 900)
await _check(handler, _class_team("team-a"), saturation=0.92, priority="interactive")
@pytest.mark.asyncio
async def test_teamless_traffic_in_two_classes_keeps_both_classes_present(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, UserAPIKeyAuth(api_key="sk-i", metadata={"priority": "interactive"}), "interactive", 20)
await _burn(handler, UserAPIKeyAuth(api_key="sk-b", metadata={"priority": "batch"}), "batch", 20)
await _burn(handler, _class_team("team-c", "batch"), "batch", 450)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-c", "batch"), saturation=0.6, priority="batch")
_held_to_class_share(exc.value)
assert "Team share: 400, Remaining: 0," in str(exc.value.detail)
@pytest.mark.asyncio
async def test_team_counters_do_not_collide_across_models_whose_names_share_a_colon(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team("prod:victim"), "interactive", 600, model="chat")
await _burn(handler, _class_team("team-x"), "interactive", 20, model="chat:prod")
await _check(handler, _class_team("victim"), saturation=0.6, priority="interactive", model="chat:prod")
def _capped_team_key(multiple: object, limit: str = "model_tpm_limit", share: int = 100) -> UserAPIKeyAuth:
return _team_key(team_limits={limit: {"shared-model": share}, "team_model_limit_max_borrow_multiple": multiple})
@ -2494,3 +2785,174 @@ async def test_a_websocket_turn_under_the_teams_max_borrow_multiple_still_borrow
handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.5)
await _admit_then_settle(handler, _capped_team_key(2), 190)
await handler.enforce_token_limits_for_frame(user_api_key_dict=_capped_team_key(2), requested_model="shared-model")
@pytest.mark.asyncio
async def test_a_configured_class_named_default_is_not_merged_with_teams_that_have_no_priority(
monkeypatch: pytest.MonkeyPatch,
):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
monkeypatch.setattr(litellm, "priority_reservation", {"default": 0.6, "batch": 0.4})
await _burn(handler, _class_team("team-a", "default"), "default", 20)
await _burn(handler, _class_team("team-n", None), None, 20)
await _burn(handler, _class_team("team-a", "default"), "default", 700)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-a", "default"), saturation=0.74, priority="default")
assert "Priority: default, Active teams in class: 1, Active classes: 2, " in str(exc.value.detail)
@pytest.mark.asyncio
async def test_tokens_from_a_nested_call_do_not_land_on_the_outer_requests_team_counter(
monkeypatch: pytest.MonkeyPatch,
):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team("team-b"), "interactive", 20)
await _check(handler, _class_team("team-a"), saturation=0.0, priority="interactive")
from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash
stash: Final = get_or_create_request_stash()
stash.owner_litellm_call_id = "outer-call"
await handler.async_log_success_event(
kwargs={
"litellm_call_id": "judge-call",
"litellm_params": {"metadata": {"model_group": "shared-model"}},
"standard_logging_object": {
"metadata": {
"user_api_key_team_id": "team-a",
"user_api_key_auth_metadata": {"priority": "interactive"},
}
},
},
response_obj=litellm.ModelResponse(
usage=litellm.Usage(prompt_tokens=600, completion_tokens=0, total_tokens=600)
),
start_time=datetime.now(),
end_time=datetime.now(),
)
await _check(handler, _class_team("team-a"), saturation=0.62, priority="interactive")
@pytest.mark.asyncio
async def test_a_team_id_spelled_like_a_pool_marker_still_counts_as_a_team(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team(pool_member("interactive")), "interactive", 20)
await _burn(handler, _class_team("team-a"), "interactive", 600)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-a"), saturation=0.62, priority="interactive")
assert "Active teams in class: 2, Active classes: 1, " in str(exc.value.detail)
@pytest.mark.asyncio
async def test_requests_sent_while_the_share_was_not_enforced_all_count_once_the_share_grows(
monkeypatch: pytest.MonkeyPatch,
):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _check(handler, _class_team("team-b"), saturation=0.0, priority="interactive")
for _ in range(70):
await _check(handler, _class_team("team-a"), saturation=0.0, priority="interactive")
await _check(handler, _class_team("team-b", "batch"), saturation=0.0, priority="batch")
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-a"), saturation=0.6, priority="interactive")
assert "Active teams in class: 1, Active classes: 2, " in str(exc.value.detail)
@pytest.mark.asyncio
async def test_a_websocket_turn_settles_on_the_team_share_once_the_registry_recovers(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True, saturation_threshold=0.1)
await _burn(handler, _class_team("team-b"), "interactive", 20)
registry: Final = handler.active_teams
handler.active_teams = _UnavailableRegistry() # pyright: ignore[reportAttributeAccessIssue] # registry down at connect
await _check(handler, _class_team("team-a"), saturation=0.0, priority="interactive")
handler.active_teams = registry
await handler.enforce_token_limits_for_frame(
user_api_key_dict=_class_team("team-a"), requested_model="shared-model"
)
await handler.async_log_success_event(
kwargs={
"call_type": "_aresponses_websocket",
"litellm_params": {"metadata": {"model_group": "shared-model"}},
"standard_logging_object": {
"metadata": {
"user_api_key_team_id": "team-a",
"user_api_key_auth_metadata": {"priority": "interactive"},
}
},
},
response_obj=litellm.ModelResponse(
usage=litellm.Usage(prompt_tokens=550, completion_tokens=0, total_tokens=550)
),
start_time=datetime.now(),
end_time=datetime.now(),
)
with pytest.raises(ProxyRateLimitError) as exc:
await handler.enforce_token_limits_for_frame(
user_api_key_dict=_class_team("team-a"), requested_model="shared-model"
)
_held_to_class_share(exc.value)
@pytest.mark.asyncio
async def test_a_frame_checked_while_the_registry_is_down_does_not_erase_the_connections_team_share(
monkeypatch: pytest.MonkeyPatch,
):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True, saturation_threshold=0.1)
await _burn(handler, _class_team("team-b"), "interactive", 20)
await _check(handler, _class_team("team-a"), saturation=0.0, priority="interactive")
registry: Final = handler.active_teams
handler.active_teams = _UnavailableRegistry() # pyright: ignore[reportAttributeAccessIssue] # registry down mid-connection
await handler.enforce_token_limits_for_frame(
user_api_key_dict=_class_team("team-a"), requested_model="shared-model"
)
handler.active_teams = registry
await handler.async_log_success_event(
kwargs={
"call_type": "_aresponses_websocket",
"litellm_params": {"metadata": {"model_group": "shared-model"}},
"standard_logging_object": {
"metadata": {
"user_api_key_team_id": "team-a",
"user_api_key_auth_metadata": {"priority": "interactive"},
}
},
},
response_obj=litellm.ModelResponse(
usage=litellm.Usage(prompt_tokens=550, completion_tokens=0, total_tokens=550)
),
start_time=datetime.now(),
end_time=datetime.now(),
)
with pytest.raises(ProxyRateLimitError) as exc:
await handler.enforce_token_limits_for_frame(
user_api_key_dict=_class_team("team-a"), requested_model="shared-model"
)
_held_to_class_share(exc.value)
@pytest.mark.asyncio
async def test_tokens_served_by_a_fallback_model_land_on_the_teams_counter_for_that_model(
monkeypatch: pytest.MonkeyPatch,
):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
await _burn(handler, _class_team("team-b"), "interactive", 20, model="other-model")
await _burn(handler, _class_team("team-a"), "interactive", 10, model="other-model")
await _check(handler, _class_team("team-a"), saturation=0.0, priority="interactive")
await handler.async_log_success_event(
kwargs={
"litellm_params": {"metadata": {"model_group": "other-model"}},
"standard_logging_object": {
"metadata": {
"user_api_key_team_id": "team-a",
"user_api_key_auth_metadata": {"priority": "interactive"},
}
},
},
response_obj=litellm.ModelResponse(
usage=litellm.Usage(prompt_tokens=550, completion_tokens=0, total_tokens=550)
),
start_time=datetime.now(),
end_time=datetime.now(),
)
await _check(handler, _class_team("team-a"), saturation=0.62, priority="interactive")
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-a"), saturation=0.62, priority="interactive", model="other-model")
_held_to_class_share(exc.value)

157
tests/unit/proxy/hooks/test_priority_class_share.py Normal file
View file

@ -0,0 +1,157 @@
"""Share math and the active-team registry behind
``priority_reservation_settings.priority_class_share_splits_equally_among_active_teams``."""
from typing import Final
import pytest
from litellm.caching.caching import DualCache
from litellm.proxy.hooks.priority_class_share import ActiveTeamRegistry, pool_member, team_share
WEIGHTS: Final = {"interactive": 0.6, "batch": 0.4}
def test_two_interactive_and_four_batch_teams_split_the_model_sixty_forty():
members: Final = {
"i1": "interactive",
"i2": "interactive",
"b1": "batch",
"b2": "batch",
"b3": "batch",
"b4": "batch",
}
interactive: Final = team_share("interactive", members, WEIGHTS, default_weight=0.2)
batch: Final = team_share("batch", members, WEIGHTS, default_weight=0.2)
assert int(500_000_000 * interactive.fraction) == 150_000_000
assert int(500_000_000 * batch.fraction) == 50_000_000
assert (interactive.active_teams_in_class, interactive.active_classes) == (2, 2)
def test_only_active_class_takes_the_whole_model():
members: Final = {"b1": "batch", "b2": "batch", "b3": "batch", "b4": "batch"}
assert int(500_000_000 * team_share("batch", members, WEIGHTS, default_weight=0.2).fraction) == 125_000_000
def test_a_single_active_team_gets_everything():
share: Final = team_share("interactive", {"i1": "interactive"}, WEIGHTS, default_weight=0.2)
assert share.fraction == 1.0
assert (share.active_teams_in_class, share.active_classes) == (1, 1)
def test_a_class_without_a_configured_weight_uses_the_default_weight():
members: Final = {"i1": "interactive", "d1": "default"}
assert team_share("default", members, WEIGHTS, default_weight=0.2).fraction == pytest.approx(0.2 / 0.8)
assert team_share("interactive", members, WEIGHTS, default_weight=0.2).fraction == pytest.approx(0.6 / 0.8)
def test_active_team_fractions_always_sum_to_one():
members: Final = {"i1": "interactive", "i2": "interactive", "b1": "batch", "d1": "default", "d2": "default"}
total: Final = sum(team_share(cls, members, WEIGHTS, default_weight=0.2).fraction for cls in members.values())
assert total == pytest.approx(1.0)
def test_zero_weights_split_the_model_equally_between_classes():
members: Final = {"a": "interactive", "b": "batch"}
assert team_share("batch", members, {"interactive": 0.0, "batch": 0.0}, default_weight=0.0).fraction == 0.5
class _Clock:
def __init__(self) -> None:
self.now = 1_000.0
def __call__(self) -> float:
return self.now
@pytest.mark.asyncio
async def test_registry_lists_every_team_seen_in_the_window_oldest_first():
clock: Final = _Clock()
registry: Final = ActiveTeamRegistry(DualCache(), time_provider=clock)
assert dict(await registry.touch("m", "a", "interactive", window=60)) == {"a": "interactive"}
clock.now += 10
assert dict(await registry.touch("m", "b", "batch", window=60)) == {"a": "interactive", "b": "batch"}
@pytest.mark.asyncio
async def test_registry_forgets_a_team_once_its_last_request_leaves_the_window():
clock: Final = _Clock()
registry: Final = ActiveTeamRegistry(DualCache(), time_provider=clock)
await registry.touch("m", "a", "interactive", window=60)
clock.now += 59
assert "a" in await registry.touch("m", "b", "batch", window=60)
clock.now += 1
assert dict(await registry.touch("m", "b", "batch", window=60)) == {"b": "batch"}
@pytest.mark.asyncio
async def test_registry_refreshes_a_teams_last_seen_time_on_every_request():
clock: Final = _Clock()
registry: Final = ActiveTeamRegistry(DualCache(), time_provider=clock)
await registry.touch("m", "a", "interactive", window=60)
clock.now += 40
await registry.touch("m", "a", "interactive", window=60)
clock.now += 40
assert "a" in await registry.touch("m", "b", "batch", window=60)
@pytest.mark.asyncio
async def test_registry_keeps_models_apart():
registry: Final = ActiveTeamRegistry(DualCache(), time_provider=_Clock())
await registry.touch("m1", "a", "interactive", window=60)
assert dict(await registry.touch("m2", "b", "batch", window=60)) == {"b": "batch"}
@pytest.mark.asyncio
async def test_registry_keeps_a_teams_newest_class_when_it_changed_class_mid_window():
clock: Final = _Clock()
registry: Final = ActiveTeamRegistry(DualCache(), time_provider=clock)
await registry.touch("m", "a", "interactive", window=60)
clock.now += 1
assert dict(await registry.touch("m", "a", "batch", window=60)) == {"a": "batch"}
@pytest.mark.asyncio
async def test_registry_lists_the_caller_under_the_class_it_just_sent_with_even_when_another_worker_wrote_later():
clock: Final = _Clock()
registry: Final = ActiveTeamRegistry(DualCache(), time_provider=clock)
clock.now += 5
await registry.touch("m", "a", "batch", window=60)
clock.now -= 5
members: Final = await registry.touch("m", "a", "interactive", window=60)
assert dict(members) == {"a": "interactive"}
assert team_share("interactive", members, WEIGHTS, 1.0).fraction == 1.0
@pytest.mark.asyncio
async def test_registry_keeps_a_team_id_that_contains_the_separator_whole():
registry: Final = ActiveTeamRegistry(DualCache(), time_provider=_Clock())
members: Final = await registry.touch("m", "owned\x1fextra", "batch", window=60)
assert dict(members) == {"owned\x1fextra": "batch"}
assert team_share("batch", members, {"interactive": 0.6, "batch": 0.4}, 1.0).fraction == 1.0
def test_a_pool_member_keeps_its_class_present_without_taking_a_team_slice():
members: Final = {pool_member("interactive"): "interactive", "a": "interactive", "c": "batch"}
share: Final = team_share("interactive", members, {"interactive": 0.6, "batch": 0.4}, 1.0)
assert share.fraction == 0.6
assert share.active_teams_in_class == 1
assert share.active_classes == 2
@pytest.mark.asyncio
async def test_registry_keeps_pool_markers_of_different_classes_apart():
registry: Final = ActiveTeamRegistry(DualCache(), time_provider=_Clock())
await registry.touch("m", pool_member("interactive"), "interactive", window=60)
members: Final = await registry.touch("m", pool_member("batch"), "batch", window=60)
assert set(members.values()) == {"interactive", "batch"}
@pytest.mark.asyncio
async def test_registry_keeps_a_team_that_was_refreshed_after_the_first_window_opened():
clock: Final = _Clock()
registry: Final = ActiveTeamRegistry(DualCache(), time_provider=clock)
await registry.touch("m", "a", "interactive", window=60)
clock.now += 50
await registry.touch("m", "a", "interactive", window=60)
clock.now += 50
assert dict(await registry.touch("m", "b", "batch", window=60)) == {"a": "interactive", "b": "batch"}