From 7bdccd7371264a32ab2481fdc7c448ffd01a79e0 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:11:57 -0700 Subject: [PATCH] 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 --- .../hooks/parallel_request_limiter_v3.py | 84 ++++++++- .../spend/test_tag_budget_enforcement.py | 160 +++++++++++++++++- .../hooks/test_parallel_request_limiter_v3.py | 153 +++++++++++++++++ 3 files changed, 390 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 17cf7382246..79de5e26a6b 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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 = ( diff --git a/tests/integration/spend/test_tag_budget_enforcement.py b/tests/integration/spend/test_tag_budget_enforcement.py index f29ee6c07c8..14470d57359 100644 --- a/tests/integration/spend/test_tag_budget_enforcement.py +++ b/tests/integration/spend/test_tag_budget_enforcement.py @@ -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: diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 8c0dcd3383c..6cdd6a81bc7 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -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 /