mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
feat(proxy): enforce tpm_limit and rpm_limit set on tag objects (#41807)
* feat(proxy): enforce tpm_limit and rpm_limit set on tag objects Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): keep tag rate limit helpers within type discipline budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): drop descriptive docstrings from tag rate limit helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover tag object rpm and tpm limits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): type the fake tag batch helper parameters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): name the over-limit tag in 429 errors Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover tag rpm limit shared across teams, orgs and users Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): format the tag descriptor match in the v3 limiter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): drop Final annotation inside loop for pyright 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> Co-authored-by: kerry <kerry@berri.ai>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
67d7ac58cd
commit
7bdccd7371
3 changed files with 390 additions and 7 deletions
|
|
@ -472,11 +472,12 @@ class RateLimitStatus(TypedDict):
|
|||
limit_remaining: int
|
||||
rate_limit_type: Literal["requests", "tokens", "max_parallel_requests"]
|
||||
descriptor_key: str
|
||||
# Only populated by the atomic_check_and_increment_by_n path. A caller
|
||||
# matching a status back to its descriptor must key on (descriptor_key,
|
||||
# descriptor_value) when this is present, not descriptor_key alone --
|
||||
# e.g. a batch charging several models' project ITPM/OTPM in one call
|
||||
# produces multiple statuses sharing the same descriptor_key.
|
||||
# Populated by the atomic_check_and_increment_by_n and windowed
|
||||
# sliding-window paths. A caller matching a status back to its
|
||||
# descriptor must key on (descriptor_key, descriptor_value) when this
|
||||
# is present, not descriptor_key alone -- e.g. a batch charging several
|
||||
# models' project ITPM/OTPM in one call, or a request carrying multiple
|
||||
# rate-limited tags, produces statuses sharing the same descriptor_key.
|
||||
descriptor_value: NotRequired[ReadOnly[str]]
|
||||
|
||||
|
||||
|
|
@ -506,6 +507,7 @@ class WindowKeyMetadata(TypedDict):
|
|||
tokens_limit: int | None
|
||||
window_size: int
|
||||
descriptor_key: str
|
||||
descriptor_value: ReadOnly[str]
|
||||
|
||||
|
||||
class AtomicCounterMeta(TypedDict):
|
||||
|
|
@ -582,6 +584,47 @@ class RequestRateLimiterStash:
|
|||
batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None
|
||||
batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = ()
|
||||
reservation_released: bool = False
|
||||
tpm_limited_tags: frozenset[str] = field(default_factory=frozenset)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TagRateLimit:
|
||||
rpm_limit: int | None
|
||||
tpm_limit: int | None
|
||||
|
||||
|
||||
class TagRateLimitResolver(Protocol):
|
||||
def __call__(self, tag_names: Sequence[str], /) -> Awaitable[Mapping[str, TagRateLimit]]: ...
|
||||
|
||||
|
||||
def _tag_rate_limit_descriptor(tag: str, limit: TagRateLimit, window_size: int) -> RateLimitDescriptor:
|
||||
rate_limit: Final[RateLimitDescriptorRateLimitObject] = {
|
||||
"requests_per_unit": limit.rpm_limit,
|
||||
"tokens_per_unit": limit.tpm_limit,
|
||||
"window_size": window_size,
|
||||
}
|
||||
return RateLimitDescriptor(key="tag", value=tag, rate_limit=rate_limit)
|
||||
|
||||
|
||||
async def resolve_tag_rate_limits_from_db(tag_names: Sequence[str]) -> Mapping[str, TagRateLimit]:
|
||||
from litellm.proxy.auth.auth_checks import get_tag_objects_batch
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
if prisma_client is None or not tag_names:
|
||||
return MappingProxyType({})
|
||||
tag_objects: Final = await get_tag_objects_batch(
|
||||
tag_names=tag_names,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
tag_name: TagRateLimit(rpm_limit=budget.rpm_limit, tpm_limit=budget.tpm_limit)
|
||||
for tag_name, tag_object in tag_objects.items()
|
||||
if (budget := tag_object.litellm_budget_table) is not None
|
||||
and (budget.rpm_limit is not None or budget.tpm_limit is not None)
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_request_stash: Final[ContextVar[RequestRateLimiterStash | None]] = ContextVar(
|
||||
|
|
@ -647,10 +690,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
self,
|
||||
internal_usage_cache: InternalUsageCache,
|
||||
time_provider: Callable[[], datetime] | None = None,
|
||||
tag_rate_limit_resolver: TagRateLimitResolver = resolve_tag_rate_limits_from_db,
|
||||
model_group_resolver: Callable[[str], str | None] = _resolve_model_group_alias_via_proxy_router,
|
||||
):
|
||||
self.internal_usage_cache = internal_usage_cache
|
||||
self._time_provider = time_provider or datetime.now
|
||||
self._tag_rate_limit_resolver = tag_rate_limit_resolver
|
||||
self._model_group_resolver = model_group_resolver
|
||||
if self.internal_usage_cache.dual_cache.redis_cache is not None:
|
||||
self.batch_rate_limiter_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
|
||||
|
|
@ -1185,6 +1230,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
"limit_remaining": limit_remaining,
|
||||
"rate_limit_type": rate_limit_type,
|
||||
"descriptor_key": key_metadata[window_key]["descriptor_key"],
|
||||
"descriptor_value": key_metadata[window_key]["descriptor_value"],
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -1489,6 +1535,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
"tokens_limit": int(tokens_limit) if tokens_limit is not None else None,
|
||||
"window_size": int(window_size),
|
||||
"descriptor_key": descriptor_key,
|
||||
"descriptor_value": descriptor_value,
|
||||
}
|
||||
return keys_to_fetch, key_metadata, gauges
|
||||
|
||||
|
|
@ -2734,6 +2781,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
|
||||
return descriptors
|
||||
|
||||
async def _create_tag_rate_limit_descriptors(self, data: Mapping[str, object]) -> tuple[RateLimitDescriptor, ...]:
|
||||
tags: Final = tuple(dict.fromkeys(get_tags_from_request_body(data)))
|
||||
if not tags:
|
||||
return ()
|
||||
tag_limits: Final = await self._tag_rate_limit_resolver(tags)
|
||||
return tuple(
|
||||
_tag_rate_limit_descriptor(tag, limit, self.window_size)
|
||||
for tag in tags
|
||||
if (limit := tag_limits.get(tag)) is not None
|
||||
)
|
||||
|
||||
def _create_rate_limit_descriptors(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -3110,7 +3168,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
if status["code"] == "OVER_LIMIT":
|
||||
descriptor_key = status["descriptor_key"]
|
||||
matching_descriptor = next(
|
||||
(desc for desc in descriptors if desc["key"] == descriptor_key),
|
||||
(
|
||||
desc
|
||||
for desc in descriptors
|
||||
if desc["key"] == descriptor_key
|
||||
and ((status_value := status.get("descriptor_value")) is None or desc["value"] == status_value)
|
||||
),
|
||||
None,
|
||||
)
|
||||
descriptor_value = matching_descriptor["value"] if matching_descriptor is not None else "unknown"
|
||||
|
|
@ -3519,6 +3582,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
return [ # mutable-ok: the shared generation reservation helpers require a list
|
||||
*descriptors,
|
||||
*self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model),
|
||||
*await self._create_tag_rate_limit_descriptors(data),
|
||||
]
|
||||
|
||||
async def _release_request_capacity_when_admitted(
|
||||
|
|
@ -3613,6 +3677,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
data=request_data,
|
||||
call_type=call_type,
|
||||
)
|
||||
stash.tpm_limited_tags = frozenset(
|
||||
d["value"]
|
||||
for d in descriptors
|
||||
if d["key"] == "tag" and d["rate_limit"] is not None and d["rate_limit"].get("tokens_per_unit") is not None
|
||||
)
|
||||
|
||||
# Only check rate limits if we have descriptors with actual limits
|
||||
if descriptors:
|
||||
|
|
@ -4299,6 +4368,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
standard_logging_metadata: dict[str, Any],
|
||||
kwargs: object,
|
||||
model_group: str | None,
|
||||
tpm_limited_tags: Set[str] = frozenset(),
|
||||
) -> list[tuple[str, str]]:
|
||||
"""
|
||||
Enumerate every (scope_key, scope_value) pair that *might* carry a
|
||||
|
|
@ -4354,6 +4424,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
targets.append(("agent", agent_id))
|
||||
if session_id:
|
||||
targets.append(("agent_session", f"{agent_id}:{session_id}"))
|
||||
targets.extend(("tag", tag) for tag in sorted(tpm_limited_tags))
|
||||
return targets
|
||||
|
||||
def _build_reservation_aware_tpm_ops(
|
||||
|
|
@ -4528,6 +4599,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
targets: Final = self._collect_tpm_scope_targets(
|
||||
standard_logging_metadata=standard_logging_metadata,
|
||||
kwargs=kwargs,
|
||||
tpm_limited_tags=stash.tpm_limited_tags if stash is not None else frozenset(),
|
||||
model_group=reconcile_model.group if reconcile_model is not None else None,
|
||||
)
|
||||
charged_targets: Final = (
|
||||
|
|
|
|||
|
|
@ -1,9 +1,12 @@
|
|||
import os
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
from integration._support.client import Gateway, eventually
|
||||
import httpx
|
||||
from integration._support.client import Gateway, eventually, string_value
|
||||
from integration._support.process import owned_proxy
|
||||
from redis import Redis
|
||||
|
||||
|
||||
def test_spend_over_a_tag_max_budget_rejects_the_next_request(gateway: Gateway) -> None:
|
||||
|
|
@ -62,6 +65,161 @@ def test_spend_over_a_tag_max_budget_rejects_the_next_request(gateway: Gateway)
|
|||
assert control.status_code == 200, control.text
|
||||
|
||||
|
||||
def test_tag_object_rpm_limit_rejects_the_second_request_across_keys(gateway: Gateway) -> None:
|
||||
tag: Final = f"tag-rpm-{uuid.uuid4().hex}"
|
||||
|
||||
def delete_tag() -> None:
|
||||
gateway.post("/tag/delete", {"name": tag})
|
||||
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.01, output_cost_per_token=0.01)
|
||||
key_a: Final = scenario.key()
|
||||
key_b: Final = scenario.key()
|
||||
gateway.post("/tag/new", {"name": tag, "rpm_limit": 1})
|
||||
scenario.cleanups.callback(delete_tag)
|
||||
first: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"tag rpm {tag}"}],
|
||||
"metadata": {"tags": [tag]},
|
||||
},
|
||||
key=key_a,
|
||||
)
|
||||
assert first.status_code == 200, first.text
|
||||
second: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"tag rpm {tag}"}],
|
||||
"metadata": {"tags": [tag]},
|
||||
},
|
||||
key=key_b,
|
||||
)
|
||||
assert second.status_code == 429, second.text
|
||||
assert "tag" in second.text.lower(), second.text
|
||||
control: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"other tag rpm {tag}"}],
|
||||
"metadata": {"tags": [f"other-{tag}"]},
|
||||
},
|
||||
key=key_b,
|
||||
)
|
||||
assert control.status_code == 200, control.text
|
||||
|
||||
|
||||
def test_tag_object_rpm_limit_is_shared_across_teams_organizations_and_users(gateway: Gateway) -> None:
|
||||
tag: Final = f"tag-rpm-{uuid.uuid4().hex}"
|
||||
|
||||
def delete_tag() -> None:
|
||||
gateway.post("/tag/delete", {"name": tag})
|
||||
|
||||
def delete_organization(identity: str) -> None:
|
||||
deleted: Final = gateway.request("DELETE", "/organization/delete", {"organization_ids": [identity]})
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.01, output_cost_per_token=0.01)
|
||||
org_a: Final = gateway.post("/organization/new", {"organization_alias": f"integration-{uuid.uuid4().hex}"})
|
||||
org_b: Final = gateway.post("/organization/new", {"organization_alias": f"integration-{uuid.uuid4().hex}"})
|
||||
org_a_id: Final = string_value(org_a["organization_id"])
|
||||
org_b_id: Final = string_value(org_b["organization_id"])
|
||||
scenario.cleanups.callback(delete_organization, org_a_id)
|
||||
scenario.cleanups.callback(delete_organization, org_b_id)
|
||||
team_a: Final = scenario.team(organization_id=org_a_id)
|
||||
team_b: Final = scenario.team(organization_id=org_b_id)
|
||||
user_1: Final = scenario.user()
|
||||
user_2: Final = scenario.user()
|
||||
key_team_a: Final = scenario.key(team_id=team_a)
|
||||
key_team_b: Final = scenario.key(team_id=team_b)
|
||||
key_user_1: Final = scenario.key(user_id=user_1)
|
||||
key_user_2: Final = scenario.key(user_id=user_2)
|
||||
gateway.post("/tag/new", {"name": tag, "rpm_limit": 3})
|
||||
scenario.cleanups.callback(delete_tag)
|
||||
|
||||
def tagged_request(key: str, request_tag: str) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"tag rpm {request_tag}"}],
|
||||
"metadata": {"tags": [request_tag]},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
for scoped_key in (key_team_a, key_team_b, key_user_1):
|
||||
admitted: Final = tagged_request(scoped_key, tag)
|
||||
assert admitted.status_code == 200, admitted.text
|
||||
blocked: Final = tagged_request(key_user_2, tag)
|
||||
assert blocked.status_code == 429, blocked.text
|
||||
assert "tag" in blocked.text.lower(), blocked.text
|
||||
control: Final = tagged_request(key_user_2, f"other-{tag}")
|
||||
assert control.status_code == 200, control.text
|
||||
|
||||
|
||||
def test_tag_object_tpm_limit_rejects_the_next_request_once_tokens_are_charged(gateway: Gateway) -> None:
|
||||
tag: Final = f"tag-tpm-{uuid.uuid4().hex}"
|
||||
|
||||
def delete_tag() -> None:
|
||||
gateway.post("/tag/delete", {"name": tag})
|
||||
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.01, output_cost_per_token=0.01)
|
||||
key: Final = scenario.key()
|
||||
gateway.post("/tag/new", {"name": tag, "tpm_limit": 39})
|
||||
scenario.cleanups.callback(delete_tag)
|
||||
first: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"tag tpm {tag}"}],
|
||||
"metadata": {"tags": [tag]},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert first.status_code == 200, first.text
|
||||
|
||||
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
|
||||
charged: Final = eventually(
|
||||
lambda: int(cache.get(f"{{tag:{tag}}}:tokens") or 0),
|
||||
lambda tokens: tokens >= 40,
|
||||
seconds=70,
|
||||
)
|
||||
assert charged >= 40, charged
|
||||
second: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"tag tpm {tag}"}],
|
||||
"metadata": {"tags": [tag]},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert second.status_code == 429, second.text
|
||||
assert "tag" in second.text.lower(), second.text
|
||||
assert "token" in second.text.lower(), second.text
|
||||
control: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"other tag tpm {tag}"}],
|
||||
"metadata": {"tags": [f"other-{tag}"]},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert control.status_code == 200, control.text
|
||||
|
||||
|
||||
def test_key_tag_rpm_limit_rejects_the_second_request_carrying_that_tag(gateway: Gateway) -> None:
|
||||
tag: Final = f"tag-rpm-{uuid.uuid4().hex}"
|
||||
with gateway.scenario() as scenario:
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
|||
RateLimitedModel,
|
||||
RateLimitResponse,
|
||||
RequestRateLimiterStash,
|
||||
TagRateLimit,
|
||||
_request_stash,
|
||||
get_or_create_request_stash,
|
||||
get_request_stash,
|
||||
|
|
@ -5006,6 +5007,158 @@ async def test_per_tag_untagged_request_governed_by_key_limit_v3(monkeypatch):
|
|||
assert "tag_per_key" not in str(exc_info.value.detail)
|
||||
|
||||
|
||||
def _static_tag_limits(limits: dict[str, TagRateLimit]):
|
||||
calls: list[tuple[str, ...]] = []
|
||||
|
||||
async def resolver(tag_names: Sequence[str]):
|
||||
calls.append(tuple(tag_names))
|
||||
return {name: limits[name] for name in tag_names if name in limits}
|
||||
|
||||
return resolver, calls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tag_object_rpm_limit_enforced_v3(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
|
||||
_request_stash.set(None)
|
||||
resolver, calls = _static_tag_limits({"cell-1": TagRateLimit(rpm_limit=2, tpm_limit=None)})
|
||||
local_cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache),
|
||||
tag_rate_limit_resolver=resolver,
|
||||
)
|
||||
|
||||
async def call(api_key: str, tags: list[str]) -> None:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key=hash_token(api_key)),
|
||||
cache=local_cache,
|
||||
data={"model": "gpt-3.5-turbo", "metadata": {"tags": tags}},
|
||||
call_type="",
|
||||
)
|
||||
|
||||
await call("sk-a", ["cell-1"])
|
||||
await call("sk-b", ["cell-1", "cell-2"])
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await call("sk-a", ["cell-1"])
|
||||
assert exc_info.value.status_code == 429
|
||||
assert "tag" in str(exc_info.value.detail)
|
||||
|
||||
for _ in range(3):
|
||||
await call("sk-a", ["cell-2"])
|
||||
await call("sk-a", [])
|
||||
assert calls == [("cell-1",), ("cell-1", "cell-2"), ("cell-1",), ("cell-2",), ("cell-2",), ("cell-2",)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tag_429_names_the_tag_that_is_over_its_limit(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
|
||||
_request_stash.set(None)
|
||||
resolver, _ = _static_tag_limits(
|
||||
{
|
||||
"cell-ok": TagRateLimit(rpm_limit=100, tpm_limit=None),
|
||||
"cell-blocked": TagRateLimit(rpm_limit=1, tpm_limit=None),
|
||||
}
|
||||
)
|
||||
local_cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache),
|
||||
tag_rate_limit_resolver=resolver,
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tag-order"))
|
||||
|
||||
async def call(tags: list[str]) -> None:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=local_cache,
|
||||
data={"model": "gpt-3.5-turbo", "metadata": {"tags": tags}},
|
||||
call_type="",
|
||||
)
|
||||
|
||||
await call(["cell-blocked"])
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await call(["cell-ok", "cell-blocked"])
|
||||
assert exc_info.value.status_code == 429
|
||||
assert "cell-blocked" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tag_object_tpm_limit_enforced_v3(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
|
||||
monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false")
|
||||
_request_stash.set(None)
|
||||
resolver, _ = _static_tag_limits({"cell-1": TagRateLimit(rpm_limit=None, tpm_limit=100)})
|
||||
local_cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache),
|
||||
tag_rate_limit_resolver=resolver,
|
||||
)
|
||||
monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total")
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tag-tpm"))
|
||||
|
||||
async def call(tags: list[str]) -> None:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=local_cache,
|
||||
data={"model": "gpt-3.5-turbo", "metadata": {"tags": tags}},
|
||||
call_type="",
|
||||
)
|
||||
|
||||
await call(["cell-1"])
|
||||
tokens_before_success = await local_cache.async_get_cache("{tag:cell-1}:tokens") or 0
|
||||
await handler.async_log_success_event(
|
||||
kwargs={
|
||||
"standard_logging_object": {"metadata": {"user_api_key_hash": user_api_key_dict.api_key}},
|
||||
"model": "gpt-3.5-turbo",
|
||||
},
|
||||
response_obj=ModelResponse(usage=Usage(prompt_tokens=60, completion_tokens=60, total_tokens=120)),
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
assert await local_cache.async_get_cache("{tag:cell-1}:tokens") == tokens_before_success + 120
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await call(["cell-1"])
|
||||
assert exc_info.value.status_code == 429
|
||||
await call([])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_tag_rate_limits_from_db_reads_budget_row(monkeypatch):
|
||||
from litellm.models.budget import LiteLLM_BudgetTable
|
||||
from litellm.models.tag import LiteLLM_TagTable
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.auth import auth_checks
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import resolve_tag_rate_limits_from_db
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
async def fake_batch(
|
||||
tag_names: Sequence[str],
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: DualCache,
|
||||
) -> dict[str, LiteLLM_TagTable]:
|
||||
return {
|
||||
"limited": LiteLLM_TagTable(
|
||||
tag_name="limited",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(rpm_limit=3, tpm_limit=None),
|
||||
),
|
||||
"spend-only": LiteLLM_TagTable(
|
||||
tag_name="spend-only",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=5.0),
|
||||
),
|
||||
"bare": LiteLLM_TagTable(tag_name="bare"),
|
||||
}
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", object())
|
||||
monkeypatch.setattr(auth_checks, "get_tag_objects_batch", fake_batch)
|
||||
|
||||
assert dict(await resolve_tag_rate_limits_from_db(["limited", "spend-only", "bare"])) == {
|
||||
"limited": TagRateLimit(rpm_limit=3, tpm_limit=None)
|
||||
}
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
assert dict(await resolve_tag_rate_limits_from_db(["limited"])) == {}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Streaming success logging mirrors x-ratelimit-* remaining values into
|
||||
# standard_logging_object.hidden_params.additional_headers so Prometheus /
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue