diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b04ef5f7087..486e3cf88a4 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1341,7 +1341,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): else: limit_value = rate_limit.get("tokens_per_unit") inc_amount = int(increment_amounts.get("tokens", 0) or 0) - if limit_value is None or inc_amount <= 0: + if limit_value is None: continue counter_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, rlt) # Counter-key TTL and window_size are conceptually distinct diff --git a/tests/e2e/models.py b/tests/e2e/models.py index af695acaa5e..f1c0ede0e85 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -43,6 +43,7 @@ class KeyLoggingCallback(BaseModel): class KeyMetadata(BaseModel): logging: list[KeyLoggingCallback] | None = None + priority: str | None = None class ObjectPermission(BaseModel): @@ -97,6 +98,7 @@ class LiteLLMBudgetTable(BaseModel): class KeyInfo(BaseModel): key_alias: str | None = None + metadata: KeyMetadata | None = None models: list[str] = [] tpm_limit: int | None = None rpm_limit: int | None = None @@ -694,6 +696,7 @@ class LiteLLMParamsBody(BaseModel): complexity_router_config: dict[str, object] | None = None mock_response: str | None = None timeout: float | None = None + tpm: int | None = None ModelMode = Literal["batch", "realtime", "image_generation"] diff --git a/tests/e2e/quota_management/ratelimit/test_dynamic_rate_limit_priority_e2e.py b/tests/e2e/quota_management/ratelimit/test_dynamic_rate_limit_priority_e2e.py new file mode 100644 index 00000000000..06b4aa1a0b5 --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/test_dynamic_rate_limit_priority_e2e.py @@ -0,0 +1,242 @@ +"""Live e2e: the v3 dynamic rate limiter's saturation-aware priority reservation. + +Covers quota_management.ratelimit.priority_generous / priority_strict: with +`dynamic_rate_limiter_v3` enabled, a model's TPM capacity is split into priority +reservations, but a reservation is only enforced once the model is saturated. + +- Generous mode (recorded usage below the saturation threshold): a key whose + priority reserves 25% of capacity keeps serving past its reservation, + borrowing the idle capacity (priority_generous.picks_under_tpm) +- Strict mode (recorded usage at/over the threshold): the over-reservation key + is blocked with the priority-flavored 429 while a key of a different priority, + still inside its own reservation, is served (priority_strict.picks_under_tpm) + +The proxy under test must run with this config (and LITELLM_LICENSE set, since +priority reservation is a premium feature): + + litellm_settings: + callbacks: ["dynamic_rate_limiter_v3"] + priority_reservation: + prod: 0.5 + dev: 0.25 + priority_reservation_settings: + saturation_threshold: 0.5 + saturation_check_cache_ttl: 1 + +The constants below mirror those values; if the proxy runs different ones the +tests fail with a message naming the required config rather than skipping. + +The limiter counts a request against the model-wide window pre-call, but tokens +only land on the counters after each response completes (there is no pre-call +token reservation at the model level), so recorded saturation always trails the +traffic that produced it. The tests therefore drive spend by summing each +body's usage.total_tokens (the counter can never be ahead of that sum) and poll +for the strict-mode block instead of expecting it on an exact call. Each test +creates its own /model/new deployment so its 60s rate-limit window and counters +are isolated from concurrent runs. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass + +import pytest +from pydantic import BaseModel, ConfigDict, ValidationError + +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call +from lifecycle import ResourceManager +from models import KeyGenerateBody, KeyMetadata, LiteLLMParamsBody +from quota_client import QuotaClient + +pytestmark = pytest.mark.e2e + +BACKEND = "anthropic/claude-haiku-4-5-20251001" +MODEL_TPM = 400 +DEV_PRIORITY = "dev" +PROD_PRIORITY = "prod" +DEV_RESERVED_TOKENS = int(MODEL_TPM * 0.25) +SATURATION_TOKENS = int(MODEL_TPM * 0.5) +CHAT_MAX_TOKENS = 16 +WINDOW_SECONDS = 60 +WINDOW_MARGIN_SECONDS = 10 +STRICT_POLL_SPEND_CEILING = int(MODEL_TPM * 0.7) + +REQUIRED_CONFIG_HINT = ( + "the proxy must run litellm_settings.callbacks=['dynamic_rate_limiter_v3'] with " + "priority_reservation {prod: 0.5, dev: 0.25} and priority_reservation_settings " + "{saturation_threshold: 0.5, saturation_check_cache_ttl: 1}; see this module's docstring" +) + + +class _ChatUsage(BaseModel): + model_config = ConfigDict(extra="ignore") + + total_tokens: int + + +class _ChatBodyWithUsage(BaseModel): + model_config = ConfigDict(extra="ignore") + + usage: _ChatUsage + + +def _total_tokens(outcome: StreamingResponse) -> int: + try: + return _ChatBodyWithUsage.model_validate_json(outcome.body).usage.total_tokens + except ValidationError: + pytest.fail(f"successful chat body must report usage.total_tokens, got: {outcome.body[:300]}") + + +@dataclass(frozen=True, slots=True) +class _Fixture: + model: str + dev_key: str + prod_key: str + + +def _dynamic_limited_model(client: QuotaClient, resources: ResourceManager, label: str) -> _Fixture: + model = f"e2e-dynpri-{label}-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model=BACKEND, api_key="os.environ/ANTHROPIC_API_KEY", tpm=MODEL_TPM), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + def _priority_key(priority: str) -> str: + key = client.proxy.generate_key( + KeyGenerateBody( + models=[model], + metadata=KeyMetadata(priority=priority), + key_alias=f"e2e-dynpri-{label}-{priority}-{unique_marker()}", + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + return key + + return _Fixture(model=model, dev_key=_priority_key(DEV_PRIORITY), prod_key=_priority_key(PROD_PRIORITY)) + + +def _chat(client: QuotaClient, key: str, model: str) -> StreamingResponse: + return client.chat(key, model, f"reply with one word {unique_marker()}", max_tokens=CHAT_MAX_TOKENS) + + +@dataclass(frozen=True, slots=True) +class _FirstOk: + sent_at: float + response: StreamingResponse + + +def _first_ok(client: QuotaClient, key: str, model: str) -> _FirstOk: + """First successful call on a fresh key opens the model's rate-limit window; + `sent_at` (captured before the winning send) is a lower bound on the window + start. A fresh key may briefly 401 until the auth cache picks it up, so + retry 401s to a deadline; a 401 never reaches the limiter, so only the + successful call consumes budget.""" + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + sent_at = time.monotonic() + outcome = _chat(client, key, model) + if outcome.ok: + return _FirstOk(sent_at=sent_at, response=outcome) + if outcome.status_code != 401 or time.monotonic() >= deadline: + require_successful_call(outcome) + time.sleep(client.proxy.poll_interval) + + +def _window_guard(first: _FirstOk, spent: int) -> None: + assert time.monotonic() < first.sent_at + WINDOW_SECONDS - WINDOW_MARGIN_SECONDS, ( + f"only {spent} tokens of spend landed before the {WINDOW_SECONDS}s rate-limit window could " + "roll; this test needs every call inside one window" + ) + + +class TestDynamicRateLimitPriority: + @pytest.mark.covers( + "quota_management.ratelimit.priority_generous.picks_under_tpm", + exercised_on=["chat_completions"], + ) + def test_generous_mode_lets_priority_borrow_past_reservation( + self, client: QuotaClient, resources: ResourceManager + ) -> None: + fixture = _dynamic_limited_model(client, resources, "generous") + + info = client.proxy.key_info(fixture.dev_key) + assert info.metadata is not None and info.metadata.priority == DEV_PRIORITY, ( + f"/key/info must echo the key's priority metadata, got {info.metadata}" + ) + + first = _first_ok(client, fixture.dev_key, fixture.model) + spent = _total_tokens(first.response) + while spent <= DEV_RESERVED_TOKENS: + _window_guard(first, spent) + assert spent < SATURATION_TOKENS, ( + f"spend reached the saturation threshold ({spent} of {SATURATION_TOKENS}) before " + f"crossing the dev reservation ({DEV_RESERVED_TOKENS}); shrink per-call spend to " + "keep the borrowing claim observable" + ) + outcome = _chat(client, fixture.dev_key, fixture.model) + assert outcome.status_code != 429, ( + f"dev key was blocked at {spent} recorded tokens, under the saturation threshold " + f"({SATURATION_TOKENS} of {MODEL_TPM}); generous mode must let it borrow past its " + f"{DEV_RESERVED_TOKENS}-token reservation. If the limiter is missing entirely, " + f"{REQUIRED_CONFIG_HINT}. 429 body: {outcome.body[:300]}" + ) + require_successful_call(outcome) + spent += _total_tokens(outcome) + + assert spent > DEV_RESERVED_TOKENS + + @pytest.mark.covers( + "quota_management.ratelimit.priority_strict.picks_under_tpm", + exercised_on=["chat_completions"], + ) + def test_strict_mode_blocks_saturated_priority_but_serves_the_other( + self, client: QuotaClient, resources: ResourceManager + ) -> None: + fixture = _dynamic_limited_model(client, resources, "strict") + + prod_warmup = _first_ok(client, fixture.prod_key, fixture.model) + first = _first_ok(client, fixture.dev_key, fixture.model) + prod_spent = _total_tokens(prod_warmup.response) + dev_spent = _total_tokens(first.response) + + while prod_spent + dev_spent < SATURATION_TOKENS: + _window_guard(prod_warmup, prod_spent + dev_spent) + outcome = _chat(client, fixture.dev_key, fixture.model) + assert outcome.status_code != 429, ( + f"dev key was blocked at {prod_spent + dev_spent} recorded tokens, before the " + f"saturation threshold ({SATURATION_TOKENS} of {MODEL_TPM}); strict enforcement " + f"must not engage early. 429 body: {outcome.body[:300]}" + ) + require_successful_call(outcome) + dev_spent += _total_tokens(outcome) + + while True: + _window_guard(prod_warmup, prod_spent + dev_spent) + assert prod_spent + dev_spent < STRICT_POLL_SPEND_CEILING, ( + f"dev key was still served at {prod_spent + dev_spent} tokens, past the saturation " + f"threshold ({SATURATION_TOKENS}) and {DEV_RESERVED_TOKENS}-token dev reservation; " + f"strict priority enforcement never engaged. Check that {REQUIRED_CONFIG_HINT}" + ) + outcome = _chat(client, fixture.dev_key, fixture.model) + if outcome.status_code == 429: + assert "Priority-based rate limit exceeded" in outcome.body, ( + f"the saturated dev key must get the priority-flavored 429, got: {outcome.body[:300]}" + ) + assert outcome.headers.get("x-litellm-priority") == DEV_PRIORITY, ( + f"the 429 must attribute the blocked priority, headers: " + f"{ {k: v for k, v in outcome.headers.items() if 'litellm' in k} }" + ) + break + require_successful_call(outcome) + dev_spent += _total_tokens(outcome) + + prod_outcome = _chat(client, fixture.prod_key, fixture.model) + require_successful_call(prod_outcome) + prod_spent += _total_tokens(prod_outcome) + assert prod_spent < int(MODEL_TPM * 0.5), ( + f"the prod fairness claim needs prod spend ({prod_spent}) inside its reservation " + f"({int(MODEL_TPM * 0.5)}); shrink per-call spend" + ) diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index c8176ca6337..6c717d6f71c 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -1771,3 +1771,95 @@ async def test_priority_429_includes_model_name_and_configured_limits(): assert "Priority: prod" in error_msg, error_msg assert "Rate limit type: tokens" in error_msg, error_msg assert "Model saturation:" in error_msg, error_msg + + +@pytest.mark.asyncio +async def test_tpm_only_model_enforces_priority_and_model_capacity(): + """Regression: a model configured with ONLY tpm (no rpm) must still be + rate limited. + + The atomic check-and-increment path used to drop any counter whose + pre-call increment was zero. Token increments are always zero pre-call + (usage lands on the counters post-call), so on a TPM-only model the + limiter evaluated no counters at all: no model-wide cap, no priority + reservation, in either mode. This test drives the real pre-call -> + log-success -> pre-call flow with no limiter internals mocked. + """ + from fastapi import HTTPException + + from litellm.types.utils import ModelResponse, Usage + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"dev": 0.25, "prod": 0.5} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "tpm-only-model" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 400, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + dev_user = UserAPIKeyAuth() + dev_user.metadata = {"priority": "dev"} + prod_user = UserAPIKeyAuth() + prod_user.metadata = {"priority": "prod"} + + async def record_usage(priority: str, total_tokens: int) -> None: + await handler.async_log_success_event( + kwargs={ + "standard_logging_object": { + "metadata": {"user_api_key_auth_metadata": {"priority": priority}}, + }, + "litellm_params": {"metadata": {"model_group": model}}, + }, + response_obj=ModelResponse( + model=model, + usage=Usage(prompt_tokens=0, completion_tokens=total_tokens, total_tokens=total_tokens), + ), + start_time=None, + end_time=None, + ) + + assert ( + await handler.async_pre_call_hook( + user_api_key_dict=dev_user, cache=dual_cache, data={"model": model}, call_type="completion" + ) + is None + ) + + await record_usage("dev", 250) + + with pytest.raises(HTTPException) as dev_blocked: + await handler.async_pre_call_hook( + user_api_key_dict=dev_user, cache=dual_cache, data={"model": model}, call_type="completion" + ) + assert dev_blocked.value.status_code == 429 + assert "Priority-based rate limit exceeded" in dev_blocked.value.detail["error"] + + assert ( + await handler.async_pre_call_hook( + user_api_key_dict=prod_user, cache=dual_cache, data={"model": model}, call_type="completion" + ) + is None + ) + + await record_usage("prod", 200) + + with pytest.raises(HTTPException) as capacity_blocked: + await handler.async_pre_call_hook( + user_api_key_dict=prod_user, cache=dual_cache, data={"model": model}, call_type="completion" + ) + assert capacity_blocked.value.status_code == 429 + assert "Model capacity reached" in capacity_blocked.value.detail["error"] 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 56bfd1829b5..84e41711176 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 @@ -4998,3 +4998,60 @@ async def test_split_usage_still_respects_the_configured_limit_type(monkeypatch) token_operations = [op for op in captured_operations if op["key"].endswith(":tokens")] assert token_operations assert all(op["increment_value"] == 7 for op in token_operations) + + +@pytest.mark.asyncio +async def test_atomic_check_with_zero_increment_still_enforces_token_limit(): + """Regression: a zero token increment must still CHECK the token limit. + + The dynamic rate limiter calls atomic_check_and_increment_by_n with + {"requests": 1, "tokens": 0} because tokens land on the counter post-call. + The payload builder used to skip any counter whose increment was <= 0, so a + TPM-only descriptor produced zero counters to evaluate and the call + returned OK with empty statuses; TPM limits were never enforced at all. + """ + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor + + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + descriptor = RateLimitDescriptor( + key="model_saturation_check", + value="tpm-only-model", + rate_limit={"tokens_per_unit": 100, "window_size": 60}, + ) + zero_token_increment: Dict[str, int] = {"requests": 1, "tokens": 0} + + under_limit = await handler.atomic_check_and_increment_by_n( + descriptors=[descriptor], + increments=[zero_token_increment], + ) + assert under_limit["overall_code"] == "OK" + assert [s["rate_limit_type"] for s in under_limit["statuses"]] == ["tokens"] + + counter_key = handler.create_rate_limit_keys( + "model_saturation_check", "tpm-only-model", "tokens" + ) + await handler.async_increment_tokens_with_ttl_preservation( + pipeline_operations=[ + RedisPipelineIncrementOperation( + key=counter_key, increment_value=150, ttl=60 + ) + ], + ) + + over_limit = await handler.atomic_check_and_increment_by_n( + descriptors=[descriptor], + increments=[zero_token_increment], + ) + assert over_limit["overall_code"] == "OVER_LIMIT" + blocked = over_limit["statuses"][0] + assert blocked["rate_limit_type"] == "tokens" + assert blocked["current_limit"] == 100 + + assert ( + await handler.internal_usage_cache.async_get_cache( + key=counter_key, litellm_parent_otel_span=None, local_only=True + ) + == 150 + )