diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 972e944f729..4c5da039511 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -468,6 +468,25 @@ end return results """ +OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT: Final = """ +-- KEYS: (window_key, counter_key) pairs; ARGV: (increment, ttl_seconds) pairs at the same indexes. +-- A counter is incremented only while its scope's rate limit window is open. A pair whose window key is +-- its own counter key has no window to wait for and is always incremented. +local applied = 0 +for i = 1, #KEYS, 2 do + local counter_key = KEYS[i + 1] + if KEYS[i] == counter_key or redis.call('EXISTS', KEYS[i]) == 1 then + local ttl_seconds = tonumber(ARGV[i + 1]) + redis.call('INCRBYFLOAT', counter_key, ARGV[i]) + if ttl_seconds > 0 and redis.call('TTL', counter_key) == -1 then + redis.call('EXPIRE', counter_key, ttl_seconds) + end + applied = applied + 1 + end +end +return applied +""" + # Redis cluster slot count REDIS_CLUSTER_SLOTS: Final = 16384 REDIS_NODE_HASHTAG_NAME: Final = "all_keys" @@ -515,9 +534,6 @@ _AUDIO_BYTES_PER_TOKEN: Final = 1600 PROJECT_ITPM_DESCRIPTOR_KEY: Final = "model_per_project_itpm" PROJECT_OTPM_DESCRIPTOR_KEY: Final = "model_per_project_otpm" PTU_TEAM_DESCRIPTOR_KEY: Final = "model_per_team_ptu" -_MODEL_SCOPED_TPM_TARGETS: Final = frozenset( - {"model_per_key", "model_per_team", "model_per_organization", "model_per_project"} -) # How long an acquired slot counts toward the in-flight total before it is # considered leaked (worker crashed without any release callback firing) and # pruned. Also the longest request duration the gauge can track: a request @@ -620,6 +636,12 @@ class ReservationAwareIncrementOperation(RedisPipelineIncrementOperation): reservation_backend: NotRequired[Literal["redis", "local"]] +class OpenWindowIncrementOperation(RedisPipelineIncrementOperation): + """A token increment Redis applies only while ``open_window_key`` holds its scope's rate limit window.""" + + open_window_key: str + + class RateLimitResponseWithDescriptors(TypedDict): descriptors: list[RateLimitDescriptor] response: RateLimitResponse @@ -668,14 +690,6 @@ class CallTypeRateLimiter(Protocol): ) -> Exception | str | dict[str, object] | None: ... -@dataclass(frozen=True, slots=True) -class AdmittedRateLimitScopes: - """Every ``(scope, value)`` the pre-call hook built a rate limit descriptor for, and the model group it checked.""" - - scopes: frozenset[tuple[str, str]] - model_group: str | None - - @dataclass(slots=True) class RequestRateLimiterStash: """ @@ -723,7 +737,6 @@ class RequestRateLimiterStash: batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = () reservation_released: bool = False tpm_limited_tags: frozenset[str] = field(default_factory=frozenset) - admitted_scopes: AdmittedRateLimitScopes | None = None @dataclass(frozen=True, slots=True) @@ -823,6 +836,11 @@ def _as_window_guarded(operation: RedisPipelineIncrementOperation) -> Reservatio ) +def _open_window_key(operation: RedisPipelineIncrementOperation) -> str | None: + open_window_key: Final = operation.get("open_window_key") + return open_window_key if isinstance(open_window_key, str) else None + + def _call_id_from_callback_kwargs(kwargs: object) -> str | None: if not isinstance(kwargs, dict): return None @@ -843,6 +861,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): batch_rate_limiter_script: _AsyncLuaScript | None batch_counter_read_script: _AsyncLuaScript | None token_increment_script: _AsyncLuaScript | None + open_window_token_increment_script: _AsyncLuaScript | None check_and_increment_by_n_script: _AsyncLuaScript | None window_guarded_token_increment_script: _AsyncLuaScript | None parallel_acquire_script: _AsyncLuaScript | None @@ -878,6 +897,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.token_increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( TOKEN_INCREMENT_SCRIPT ) + self.open_window_token_increment_script = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT + ) + ) self.check_and_increment_by_n_script = ( self.internal_usage_cache.dual_cache.redis_cache.async_register_script(CHECK_AND_INCREMENT_BY_N_SCRIPT) ) @@ -899,6 +923,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.batch_rate_limiter_script = None self.batch_counter_read_script = None self.token_increment_script = None + self.open_window_token_increment_script = None self.check_and_increment_by_n_script = None self.window_guarded_token_increment_script = None self.parallel_acquire_script = None @@ -4209,7 +4234,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 ) - stash.admitted_scopes = self._admitted_scopes(descriptors, requested_model) # Only check rate limits if we have descriptors with actual limits if descriptors: @@ -4545,8 +4569,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Execute token increment script grouped by hash tag for cluster compatibility. """ - if self.token_increment_script is None: + token_increment_script: Final = self.token_increment_script + if token_increment_script is None: return + open_window_script: Final = self.open_window_token_increment_script # Group operations by hash tag for Redis cluster compatibility operation_keys: Final = [op["key"] for op in pipeline_operations] @@ -4555,6 +4581,26 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for _hash_tag, group_keys in key_groups.items(): # Get operations for this hash tag group group_operations = [op for op in pipeline_operations if op["key"] in group_keys] + window_keys = [_open_window_key(op) for op in group_operations] + if open_window_script is not None and any(window_key is not None for window_key in window_keys): + # One script for the whole group, so its plain and window-gated increments land together + await self._run_token_increment_script( + OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT, + open_window_script, + tuple( + itertools.chain.from_iterable( + (window_key or op["key"], op["key"]) + for op, window_key in zip(group_operations, window_keys, strict=True) + ) + ), + tuple( + itertools.chain.from_iterable( + (op["increment_value"], op["ttl"] or 0) for op in group_operations + ) + ), + group_operations, + ) + continue keys = [] args = [] @@ -4572,25 +4618,34 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): keys.append(op["key"]) args.extend([op["increment_value"], ttl_value]) - if self._defer_token_increment_script(keys, args, group_operations): - continue - await self.token_increment_script( - keys=keys, - args=args, + await self._run_token_increment_script( + TOKEN_INCREMENT_SCRIPT, token_increment_script, keys, args, group_operations ) + async def _run_token_increment_script( + self, + source: str, + script: _AsyncLuaScript, + keys: Sequence[str], + args: Sequence[float], + group_operations: Sequence[RedisPipelineIncrementOperation], + ) -> None: + if not self._defer_token_increment_script(source, script, keys, args, group_operations): + await script(keys=keys, args=args) + def _defer_token_increment_script( self, - keys: list[str], - args: list[int], - group_operations: list["RedisPipelineIncrementOperation"], + source: str, + script: _AsyncLuaScript, + keys: Sequence[str], + args: Sequence[float], + group_operations: Sequence[RedisPipelineIncrementOperation], ) -> bool: """Declared into the request's post-call pipeline instead of its own EVALSHA round trip; a failed script falls back to the plain increment pipeline for its own group, as the direct path does.""" redis_cache: Final = self.internal_usage_cache.dual_cache.redis_cache - script: Final = self.token_increment_script batch: Final = None if redis_cache is None else active_post_call_redis_batch(redis_cache) - if batch is None or script is None: + if batch is None: return False async def fall_back(future: asyncio.Future[object]) -> None: @@ -4603,10 +4658,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): future.exception(), ) await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( - increment_list=group_operations, + increment_list=[*group_operations], ) - batch.script(TOKEN_INCREMENT_SCRIPT, script, keys, args).on_settled(fall_back) + batch.script(source, script, keys, args).on_settled(fall_back) return True @with_service_target("rate_limits") @@ -5030,36 +5085,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): targets.extend(("tag", tag) for tag in sorted(tpm_limited_tags)) return targets - def _admitted_scopes( - self, descriptors: Sequence[RateLimitDescriptor], requested_model: str | None - ) -> AdmittedRateLimitScopes: - model: Final = self._rate_limited_model(requested_model) - return AdmittedRateLimitScopes( - scopes=frozenset((d["key"], d["value"]) for d in descriptors), - model_group=model.group if model is not None else None, - ) - - @staticmethod - def _rate_limited_targets( - targets: Sequence[tuple[str, str]], - admitted: AdmittedRateLimitScopes | None, - model_group: str | None, - ) -> Sequence[tuple[str, str]]: - """ - Keep the TPM targets admission built a rate limit descriptor for, plus - the per-model targets when tokens are settled against a model group - admission did not check. Every target stays when no admission was - recorded for the call. - """ - if admitted is None: - return targets - served_admitted_group: Final = model_group == admitted.model_group - return tuple( - target - for target in targets - if target in admitted.scopes or (not served_admitted_group and target[0] in _MODEL_SCOPED_TPM_TARGETS) - ) - def _build_reservation_aware_tpm_ops( self, targets: Sequence[tuple[str, str]], @@ -5096,6 +5121,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) return ops + def _build_open_window_tpm_ops( + self, + targets: Sequence[tuple[str, str]], + tokens: int, + ) -> list[OpenWindowIncrementOperation]: + """Charge ``tokens`` to each target's counter while that scope has an open rate limit window.""" + if tokens == 0: + return [] + return [ + OpenWindowIncrementOperation( + key=self.create_rate_limit_keys(scope_key, scope_value, "tokens"), + increment_value=tokens, + ttl=self.window_size, + open_window_key=f"{{{scope_key}:{scope_value}}}:window", + ) + for scope_key, scope_value in targets + ] + def _build_project_reservation_op( self, scope: tuple[str, str], @@ -5240,15 +5283,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for scope in sorted(reserved_scopes) if scope[0] in ("agent", "agent_session") and scope not in targets ) - limited_targets: Final = self._rate_limited_targets( - targets=targets, - admitted=stash.admitted_scopes if stash is not None else None, - model_group=reconcile_model.group if reconcile_model is not None else None, - ) charged_targets: Final = ( - [target for target in limited_targets if target[0] != "model_per_team"] + [target for target in targets if target[0] != "model_per_team"] if self._key_owns_model_tpm_limit_from_request_metadata(request_metadata, reconcile_model) - else limited_targets + else targets ) if reserved_tokens > 0 and total_tokens < reserved_tokens: verbose_proxy_logger.debug( @@ -5259,12 +5297,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) pipeline_operations.extend( self._build_reservation_aware_tpm_ops( - targets=charged_targets, + targets=[target for target in charged_targets if target in reserved_scopes], reserved_scopes=reserved_scopes, actual_tokens=total_tokens, reserved_tokens=reserved_tokens, ) ) + pipeline_operations.extend( + self._build_open_window_tpm_ops( + targets=[target for target in charged_targets if target not in reserved_scopes], + tokens=total_tokens, + ) + ) pipeline_operations.extend( self._build_team_ptu_tpm_ops( standard_logging_metadata, # pyright: ignore[reportUnknownArgumentType] # untyped logging metadata diff --git a/tests/integration/spend/test_rate_limit_token_counters.py b/tests/integration/spend/test_rate_limit_token_counters.py index d51dbe6b428..dc0d133900d 100644 --- a/tests/integration/spend/test_rate_limit_token_counters.py +++ b/tests/integration/spend/test_rate_limit_token_counters.py @@ -7,18 +7,23 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, replace from hashlib import sha256 +from pathlib import Path from typing import Final import anthropic import httpx import openai import pytest -from integration._support.client import Gateway, Scenario, eventually, held +import yaml +from integration._support.client import Gateway, Scenario, eventually, held, object_value from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue from redis import Redis +from litellm.constants import LITELLM_PROXY_ADMIN_NAME + USAGE: Final = {"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20} WIRE_TOTAL_TOKENS: Final = 8 Fields = Callable[[str], Mapping[str, JsonValue]] @@ -316,7 +321,7 @@ def test_every_endpoint_charges_only_the_rate_limited_scope(gateway: Gateway, en assert_only_charged(cache, identity.scopes, {"api_key"}, WIRE_TOTAL_TOKENS) -def test_fallback_hop_charges_the_rate_limited_model_it_served(gateway: Gateway) -> None: +def test_fallback_hop_charges_the_open_window_of_the_rate_limited_model_it_served(gateway: Gateway) -> None: prompt: Final = f"token counters fallback {uuid.uuid4().hex}" with ( wire_server(outage) as primary_upstream, @@ -332,11 +337,13 @@ def test_fallback_hop_charges_the_rate_limited_model_it_served(gateway: Gateway) key=lambda _model: {"model_tpm_limit": {fallback: 100000}}, ) identity: Final = caller(scenario, cell, fallback) + direct: Final = chat(gateway, identity, fallback, prompt + " direct") + assert direct.status_code == 200, direct.text response: Final = chat(gateway, identity, primary, prompt, fallbacks=[fallback]) assert response.status_code == 200, response.text assert response.headers["x-litellm-model-group"] == fallback assert len(primary_upstream.drain()) == 1 - assert len(fallback_upstream.drain()) == 1 + assert len(fallback_upstream.drain()) == 2 hashed: Final = sha256(identity.key.encode()).hexdigest() scopes: Final = { @@ -346,7 +353,7 @@ def test_fallback_hop_charges_the_rate_limited_model_it_served(gateway: Gateway) cache, {**scopes, "primary_model_per_key": counter("model_per_key", f"{hashed}:{primary}")}, {"model_per_key"}, - USAGE["total_tokens"], + 2 * USAGE["total_tokens"], ) @@ -420,3 +427,103 @@ def test_upstream_outage_charges_no_token_counter(gateway: Gateway) -> None: lambda values: set(values.values()) == {None}, holding=3, ) + + +def test_numeric_body_user_is_charged_to_its_rate_limited_end_user_counter(gateway: Gateway) -> None: + customer: Final = uuid.uuid4().int % 10**15 + with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + gateway.post("/customer/new", {"user_id": str(customer), "rpm_limit": 1000}) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [str(customer)]}) + identity: Final = replace( + caller(scenario, Cell("numeric_end_user", charged="end_user"), model), body={"user": customer} + ) + response: Final = chat(gateway, identity, model, f"token counters numeric user {uuid.uuid4().hex}") + assert response.status_code == 200, response.text + assert [json.loads(seen.body)["user"] for seen in upstream.drain()] == [customer] + + assert_only_charged( + cache, + {**identity.scopes, "end_user": counter("end_user", str(customer))}, + {"end_user"}, + USAGE["total_tokens"], + ) + + +USER_HEADER: Final = "x-integration-openwebui-user" + + +def proxy_admin(scenario: Scenario, tpm_limit: int) -> str: + return scenario.user(user_id=LITELLM_PROXY_ADMIN_NAME, user_role="proxy_admin", tpm_limit=tpm_limit) + + +def internal_user(scenario: Scenario, tpm_limit: int) -> str: + return scenario.user(user_role="internal_user", tpm_limit=tpm_limit) + + +def master_key_request(scenario: Scenario, _user: str, body: Mapping[str, JsonValue]) -> httpx.Response: + return scenario.gateway.request("POST", "/v1/chat/completions", body) + + +def header_mapped_service_key_request( + scenario: Scenario, user: str, body: Mapping[str, JsonValue] +) -> httpx.Response: + return scenario.gateway.request( + "POST", "/v1/chat/completions", body, key=scenario.key(), headers={USER_HEADER: user} + ) + + +def config_with_user_header_mapping(directory: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **object_value(config["general_settings"]), + "user_header_mappings": [{"header_name": USER_HEADER, "litellm_user_role": "internal_user"}], + } + path: Final = directory / "config.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.mark.parametrize( + "limited_user, unchecked_request", + [(proxy_admin, master_key_request), (internal_user, header_mapped_service_key_request)], + ids=["master_key", "user_header_mapping"], +) +def test_unchecked_traffic_counts_toward_the_user_limit_another_request_opened( + gateway: Gateway, + tmp_path: Path, + limited_user: Callable[[Scenario, int], str], + unchecked_request: Callable[[Scenario, str, Mapping[str, JsonValue]], httpx.Response], +) -> None: + redis_state: Final = {"REDIS_HOST": os.environ["REDIS_HOST"], "REDIS_PORT": os.environ["REDIS_PORT"]} + with ( + wire_server(completion) as upstream, + owned_proxy(gateway, tmp_path, redis_state, config=config_with_user_header_mapping(tmp_path)) as candidate, + candidate.scenario() as scenario, + redis_counters() as cache, + ): + model: Final = scenario.model(api_base=upstream.url + "/v1") + user: Final = limited_user(scenario, 100000) + user_counter: Final = {"user": counter("user", user)} + cache.delete(counter("user", user), f"{{user:{user}}}:window") + + checked: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"checked {uuid.uuid4().hex}"}]}, + key=scenario.key(user_id=user), + ) + assert checked.status_code == 200, checked.text + eventually(lambda: read_counters(cache, user_counter)["user"], lambda tokens: tokens == USAGE["total_tokens"]) + + unchecked: Final = unchecked_request( + scenario, user, {"model": model, "messages": [{"role": "user", "content": f"unchecked {uuid.uuid4().hex}"}]} + ) + assert unchecked.status_code == 200, unchecked.text + assert len(upstream.drain()) == 2 + + eventually( + lambda: read_counters(cache, user_counter)["user"], + lambda tokens: tokens == 2 * USAGE["total_tokens"], + seconds=30, + ) diff --git a/tests/integration/spend/test_redis_ttl_preserving_token_increment.py b/tests/integration/spend/test_redis_ttl_preserving_token_increment.py index 8b7a21f73d5..a35602bdc98 100644 --- a/tests/integration/spend/test_redis_ttl_preserving_token_increment.py +++ b/tests/integration/spend/test_redis_ttl_preserving_token_increment.py @@ -7,6 +7,7 @@ from redis import Redis from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache +from litellm.proxy.hooks.parallel_request_limiter_v3 import OpenWindowIncrementOperation from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler, ) @@ -63,3 +64,43 @@ async def test_async_increment_tokens_with_ttl_preservation() -> None: finally: await redis_cache.async_delete_cache(key_with_ttl) await redis_cache.async_delete_cache(key_without_ttl) + + +@pytest.mark.asyncio +async def test_open_window_token_increment_charges_only_scopes_whose_window_is_open() -> None: + redis_host: Final = os.environ["REDIS_HOST"] + redis_port: Final = int(os.environ["REDIS_PORT"]) + redis_cache: Final = RedisCache(host=redis_host, port=redis_port) + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis_cache)) + ) + suffix: Final = uuid.uuid4().hex[:8] + limited: Final = f"{{user:limited-{suffix}}}" + unlimited: Final = f"{{user:unlimited-{suffix}}}" + ungated: Final = f"{{team:ungated-{suffix}}}" + + with Redis(host=redis_host, port=redis_port) as raw: + try: + assert raw.set(f"{limited}:window", "1", ex=60) is True + assert raw.set(f"{limited}:tokens", "5", ex=60) is True + await handler.async_increment_tokens_with_ttl_preservation( + pipeline_operations=[ + RedisPipelineIncrementOperation(key=f"{ungated}:tokens", increment_value=20, ttl=60), + *( + OpenWindowIncrementOperation( + key=f"{scope}:tokens", increment_value=20, ttl=60, open_window_key=f"{scope}:window" + ) + for scope in (limited, unlimited) + ), + ] + ) + + assert raw.mget([f"{limited}:tokens", f"{unlimited}:tokens", f"{unlimited}:window"]) == [b"25", None, None] + assert raw.get(f"{ungated}:tokens") == b"20" + assert 0 < raw.ttl(f"{limited}:tokens") <= 60 + assert 0 < raw.ttl(f"{ungated}:tokens") <= 60 + finally: + raw.delete( + *(f"{scope}:{suffix_key}" for scope in (limited, unlimited) for suffix_key in ("window", "tokens")), + f"{ungated}:tokens", + ) diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 5de785e44db..c23f88fa891 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -30,8 +30,10 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.proxy.auth.auth_object_prefetch import AUTH_OBJECTS_TARGET from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT, PARALLEL_RELEASE_SCRIPT, TOKEN_INCREMENT_SCRIPT, + OpenWindowIncrementOperation, ParallelSlotAcquisition, RequestRateLimiterStash, PROXY_MaxParallelRequestsHandler_v3, @@ -231,6 +233,67 @@ async def test_a_failed_operation_fails_only_its_owner_and_the_owner_applies_its assert await limiter.internal_usage_cache.dual_cache.in_memory_cache.async_get_cache("{team:t1}:tokens") is None +def _open_window_op(scope: str) -> OpenWindowIncrementOperation: + return OpenWindowIncrementOperation( + key=scope + ":tokens", increment_value=10, ttl=60, open_window_key=scope + ":window" + ) + + +@pytest.mark.asyncio +async def test_a_group_with_open_window_charges_rides_the_post_call_pipeline_as_one_script(): + client = FakeClient(_ok_replies) + redis_cache = PostCallFakeRedisCache(client) + limiter = _limiter(redis_cache) + + with request_redis_batch_scope(): + await limiter.async_increment_tokens_with_ttl_preservation( + [*_token_ops("{team:t1}:tokens"), _open_window_op("{api_key:k1}")] + ) + await limiter.async_increment_tokens_with_ttl_preservation(_token_ops("{user:u1}:tokens")) + await flush_post_call_redis_batches() + + assert len(client.pipelines) == 1 + assert [command for command in client.pipelines[0].commands if command[0] == "EVALSHA"] == [ + ( + "EVALSHA", + sha_of(OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT), + 4, + "{team:t1}:tokens", + "{team:t1}:tokens", + "{api_key:k1}:window", + "{api_key:k1}:tokens", + 10, + 60, + 10, + 60, + ), + ("EVALSHA", sha_of(TOKEN_INCREMENT_SCRIPT), 1, "{user:u1}:tokens", 10, 60), + ] + assert redis_cache.alone == [] + + +@pytest.mark.asyncio +async def test_a_failed_open_window_script_falls_back_to_the_plain_increment(): + def replies(command: tuple[object, ...]) -> object: + if command[0] == "EVALSHA" and command[1] == sha_of(OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT): + return Exception("ERR Lua") + return _ok_replies(command) + + client = FakeClient(replies) + redis_cache = PostCallFakeRedisCache(client) + limiter = _limiter(redis_cache) + + with request_redis_batch_scope(): + await limiter.async_increment_tokens_with_ttl_preservation( + [*_token_ops("{team:t1}:tokens"), _open_window_op("{user:u1}")] + ) + await flush_post_call_redis_batches() + + assert redis_cache.alone == [("INCRBYFLOAT", "{team:t1}:tokens", 10), ("INCRBYFLOAT", "{user:u1}:tokens", 10)] + assert await limiter.internal_usage_cache.dual_cache.in_memory_cache.async_get_cache("{team:t1}:tokens") == 10 + assert await limiter.internal_usage_cache.dual_cache.in_memory_cache.async_get_cache("{user:u1}:tokens") == 10 + + @pytest.mark.asyncio async def test_a_failed_slot_release_script_releases_the_slot_in_memory(): def replies(command: tuple[object, ...]) -> object: diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 82b5bab8362..f4c9371f8b7 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -6,7 +6,7 @@ import asyncio import logging import sys import time -from collections.abc import Awaitable, Callable, Iterator, Sequence +from collections.abc import Iterator, Sequence from contextlib import contextmanager from datetime import datetime, timedelta, timezone from typing import Any, Dict, Final, List, Optional @@ -8688,6 +8688,7 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit( assert "expensive" in str(denied.value.detail) + _SCOPE_TOKEN_KEYS: Final = frozenset( { "{api_key:key-hash}:tokens", @@ -8700,170 +8701,69 @@ _SCOPE_TOKEN_KEYS: Final = frozenset( ) -def _scoped_success_kwargs( - model_group: str, key_metadata: Dict[str, object] | None = None, call_id: str | None = None -) -> Dict[str, object]: +def _scoped_success_kwargs(**body: object) -> dict[str, object]: return { - **({} if call_id is None else {"litellm_call_id": call_id}), "standard_logging_object": { "metadata": {"user_api_key_hash": "key-hash", "user_api_key_user_id": "u", "user_api_key_team_id": "t"} }, "litellm_params": { "metadata": { - "model_group": model_group, - "user_api_key_metadata": key_metadata or {}, + "model_group": "gpt-4o-mini", + "user_api_key_metadata": {}, "user_api_key_team_metadata": {}, } }, - "model": model_group, + "model": "gpt-4o-mini", + **body, } -async def _admit_scoped_request( - handler: _PROXY_MaxParallelRequestsHandler, auth: UserAPIKeyAuth, model: str, call_id: str | None = None -) -> None: +async def _admit_scoped_request(handler: _PROXY_MaxParallelRequestsHandler, auth: UserAPIKeyAuth) -> None: await handler.async_pre_call_hook( user_api_key_dict=auth, cache=handler.internal_usage_cache.dual_cache, - data={ - "model": model, - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 20, - **({} if call_id is None else {"litellm_call_id": call_id}), - }, + data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 20}, call_type="acompletion", ) -def _charged_token_keys(handler: _PROXY_MaxParallelRequestsHandler, kwargs: Dict[str, object]) -> frozenset[str]: +def _charged_windows(handler: _PROXY_MaxParallelRequestsHandler, kwargs: dict[str, object]) -> dict[str, object]: operations: Final = handler._build_success_event_pipeline_operations( kwargs=kwargs, response_obj=ModelResponse(usage=Usage(prompt_tokens=2, completion_tokens=1, total_tokens=3)), rate_limit_type="total", ) - return frozenset(operation["key"] for operation in operations) + return {operation["key"]: operation.get("open_window_key") for operation in operations} @pytest.mark.asyncio @pytest.mark.parametrize( - "limits, charged", + "limits, reserved", [ ({}, frozenset()), - ({"rpm_limit": 100}, frozenset({"{api_key:key-hash}:tokens"})), - ({"max_parallel_requests": 10}, frozenset({"{api_key:key-hash}:tokens"})), + ({"rpm_limit": 100}, frozenset()), + ({"user_rpm_limit": 100}, frozenset()), ({"tpm_limit": 1000}, frozenset({"{api_key:key-hash}:tokens"})), - ({"tpm_limit": 1000, "metadata": {"tpm_limit_type": "dynamic"}}, frozenset({"{api_key:key-hash}:tokens"})), - ({"user_rpm_limit": 100}, frozenset({"{user:u}:tokens"})), ({"team_tpm_limit": 1000}, frozenset({"{team:t}:tokens"})), - ({"team_member_rpm_limit": 100}, frozenset({"{team_member:t:u}:tokens"})), - ( - {"metadata": {"model_rpm_limit": {"gpt-4o-mini": 100}}}, - frozenset({"{model_per_key:key-hash:gpt-4o-mini}:tokens"}), - ), - ], - ids=[ - "no_limits", - "key_rpm", - "key_max_parallel", - "key_tpm_reserved", - "key_dynamic_tpm", - "user_rpm", - "team_tpm_reserved", - "team_member_rpm", - "key_model_rpm", ], + ids=["no_limits", "key_rpm", "user_rpm", "key_tpm_reserved", "team_tpm_reserved"], ) -async def test_success_charges_tokens_only_to_scopes_admission_rate_limited( - limits: Dict[str, object], charged: frozenset[str] +async def test_success_charges_every_unreserved_scope_only_while_its_rate_limit_window_is_open( + limits: dict[str, object], reserved: frozenset[str] ) -> None: handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) - await _admit_scoped_request( - handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", **limits), "gpt-4o-mini" - ) + await _admit_scoped_request(handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", **limits)) - assert _charged_token_keys(handler, _scoped_success_kwargs("gpt-4o-mini")) == charged - - -@pytest.mark.asyncio -async def test_fallback_success_charges_the_model_scopes_of_the_group_that_served_it() -> None: - key_metadata: Final = {"model_tpm_limit": {"fallback-model": 1000}} - handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) - await _admit_scoped_request( - handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", metadata=key_metadata), "primary-model" - ) - - assert _charged_token_keys(handler, _scoped_success_kwargs("fallback-model", key_metadata)) == { - "{model_per_key:key-hash:fallback-model}:tokens" + assert _charged_windows(handler, _scoped_success_kwargs()) == { + key: None if key in reserved else key.removesuffix(":tokens") + ":window" for key in _SCOPE_TOKEN_KEYS } - @pytest.mark.asyncio -@pytest.mark.parametrize( - "success_end_user, charged", - [ - (None, frozenset({"{end_user:admitted-customer}:tokens"})), - ("admitted-customer", frozenset({"{end_user:admitted-customer}:tokens"})), - ("other-customer", frozenset()), - ], - ids=["metadata_end_user", "same_body_user", "different_body_user"], -) -async def test_success_charges_the_end_user_admission_checked_and_no_other( - success_end_user: str | None, charged: frozenset[str] -) -> None: +async def test_success_charges_a_numeric_body_user_to_the_window_its_admission_opened() -> None: handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) await _admit_scoped_request( - handler, - UserAPIKeyAuth(api_key="key-hash", end_user_id="admitted-customer", end_user_rpm_limit=100), - "gpt-4o-mini", - ) - kwargs: Final = _scoped_success_kwargs("gpt-4o-mini") - kwargs["standard_logging_object"] = { - "metadata": {"user_api_key_hash": "key-hash", "user_api_key_end_user_id": "admitted-customer"} - } - if success_end_user is not None: - kwargs["user"] = success_end_user - - assert _charged_token_keys(handler, kwargs) == charged - - -@pytest.mark.asyncio -async def test_success_on_an_aliased_model_charges_the_model_limit_of_the_resolved_group() -> None: - key_metadata: Final = {"model_rpm_limit": {"target": 100}} - handler: Final = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()), model_group_resolver=_resolve_alias_to_target - ) - await _admit_scoped_request( - handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", metadata=key_metadata), "alias" + handler, UserAPIKeyAuth(api_key="key-hash", end_user_id="123", end_user_rpm_limit=100) ) - assert _charged_token_keys(handler, _scoped_success_kwargs("target", key_metadata)) == { - "{model_per_key:key-hash:target}:tokens" - } - -async def _no_stash(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: - _request_stash.set(None) - return None - - -async def _nested_call(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: - await _admit_scoped_request( - handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t"), "gpt-4o-mini", "proxy-request" - ) - return "nested-call" - - -async def _stash_without_admission(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: - _request_stash.set(RequestRateLimiterStash()) - return None - - -@pytest.mark.asyncio -@pytest.mark.parametrize("arrange", [_no_stash, _nested_call, _stash_without_admission]) -async def test_success_charges_every_scope_when_the_call_was_not_admitted( - arrange: Callable[[_PROXY_MaxParallelRequestsHandler], Awaitable[str | None]], -) -> None: - handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) - call_id: Final = await arrange(handler) - - assert _charged_token_keys(handler, _scoped_success_kwargs("gpt-4o-mini", call_id=call_id)) == _SCOPE_TOKEN_KEYS + assert _charged_windows(handler, _scoped_success_kwargs(user=123))["{end_user:123}:tokens"] == "{end_user:123}:window"