diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 567b8452ae2..f35659594ca 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -472,6 +472,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" @@ -622,6 +641,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 @@ -830,6 +855,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 @@ -850,6 +880,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 @@ -885,6 +916,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) ) @@ -906,6 +942,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 @@ -4589,8 +4626,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] @@ -4599,6 +4638,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 = [] @@ -4616,25 +4675,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: @@ -4647,10 +4715,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") @@ -5110,6 +5178,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], @@ -5271,12 +5357,25 @@ 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, ) ) + unreserved_targets: Final = [target for target in charged_targets if target not in reserved_scopes] + # A call that never passed admission, such as a proxy sub-call, opened no window, and the sub-call + # precheck reads its counters directly, so it charges them unconditionally + pipeline_operations.extend( + self._build_open_window_tpm_ops(targets=unreserved_targets, tokens=total_tokens) + if stash is not None + else self._build_reservation_aware_tpm_ops( + targets=unreserved_targets, + reserved_scopes=reserved_scopes, + actual_tokens=total_tokens, + reserved_tokens=0, + ) + ) pipeline_operations.extend( self._build_team_ptu_tpm_ops( standard_logging_metadata, # pyright: ignore[reportUnknownArgumentType] # untyped logging metadata diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_compaction.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_compaction.py index 38aef23c630..587fef4e2f5 100644 --- a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_compaction.py +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_compaction.py @@ -140,6 +140,20 @@ def test_caller_without_the_summary_model_sees_the_edit_error(compacting: Rig) - assert len(compacting.backend_calls(received)) == 1 +def test_caller_over_the_summary_models_tpm_limit_sees_the_edit_error(compacting: Rig) -> None: + with compacting.gateway.scenario() as scenario: + key: Final = scenario.key(model_tpm_limit={CLASSIFIER_MODEL: 1}) + within_nonce: Final = nonce() + compacting.register(within_nonce, chat_backend(tool_id(within_nonce), stream=False), summarizing_or_classifying) + within: Final = applied_edit(json_body(compacting.post(long_request(within_nonce), key=key))) + assert within["type"] == "compact_20260112" and "error" not in within, within + over_nonce: Final = nonce() + compacting.register(over_nonce, chat_backend(tool_id(over_nonce), stream=False), summarizing_or_classifying) + over: Final = json_body(compacting.post(long_request(over_nonce), key=key)) + assert applied_edit(over)["error"] == "summary_model_rate_limit_exceeded", over + assert compacting.classifier_calls(compacting.received(over_nonce)) == () + + def assert_summarized_and_classified_on_their_models(rig: Rig, request_nonce: str, body: dict[str, JsonValue]) -> str: received: Final = rig.received(request_nonce) targets: Final = [request.target for request in received] diff --git a/tests/integration/spend/test_rate_limit_token_counters.py b/tests/integration/spend/test_rate_limit_token_counters.py new file mode 100644 index 00000000000..1936b6dc260 --- /dev/null +++ b/tests/integration/spend/test_rate_limit_token_counters.py @@ -0,0 +1,533 @@ +import asyncio +import json +import os +import uuid +from collections.abc import Callable, Iterator, Mapping, Set +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 +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.redis_process import owned_redis +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]] + + +def completion(request: Request) -> Reply: + assert request.target == "/v1/chat/completions" + body: Final = json.loads(request.body) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": body["model"], + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": USAGE, + } + ).encode() + ) + + +def outage(request: Request) -> Reply: + assert request.target == "/v1/chat/completions" + return Reply(status=500, body=json.dumps({"error": {"message": "synthetic outage"}}).encode()) + + +def openai_endpoints(request: Request) -> Reply: + body: Final = json.loads(request.body) + stream: Final = body.get("stream") is True + if request.target == "/v1/chat/completions": + return chat_reply("chatcmpl-" + uuid.uuid4().hex, body["model"], "ok", stream=stream) + assert request.target == "/v1/responses", request.target + return responses_reply("resp_" + uuid.uuid4().hex, body["model"], "ok", stream=stream) + + +def counter(scope: str, value: str) -> str: + return f"{{{scope}:{value}}}:tokens" + + +def read_counters(cache: Redis, keys: Mapping[str, str]) -> Mapping[str, float | None]: + values: Final = cache.mget(list(keys.values())) + return {name: None if raw is None else float(raw) for name, raw in zip(keys, values, strict=True)} + + +def no_fields(_model: str) -> Mapping[str, JsonValue]: + return {} + + +@dataclass(frozen=True, slots=True) +class Caller: + key: str + user: str + team: str + scopes: Mapping[str, str] + body: Mapping[str, JsonValue] + + +@dataclass(frozen=True, slots=True) +class Cell: + name: str + charged: str + also_charged: frozenset[str] = frozenset() + key: Fields = no_fields + user: Fields = no_fields + team: Fields = no_fields + organization: Fields | None = None + project: Fields | None = None + customer: Fields | None = None + tag: Fields | None = None + + +def caller(scenario: Scenario, cell: Cell, model: str) -> Caller: + gateway: Final = scenario.gateway + organization: Final = None if cell.organization is None else scenario.organization(**cell.organization(model)) + team: Final = scenario.team( + **cell.team(model), **({} if organization is None else {"organization_id": organization}) + ) + user: Final = scenario.user(user_role="internal_user", **cell.user(model)) + gateway.post("/team/member_add", {"team_id": team, "member": {"role": "user", "user_id": user}}) + project: Final = None if cell.project is None else scenario.project(team, **cell.project(model)) + key: Final = scenario.key( + user_id=user, team_id=team, **({} if project is None else {"project_id": project}), **cell.key(model) + ) + hashed: Final = sha256(key.encode()).hexdigest() + scopes: Final = { + "api_key": counter("api_key", hashed), + "user": counter("user", user), + "team": counter("team", team), + "team_member": counter("team_member", f"{team}:{user}"), + "model_per_key": counter("model_per_key", f"{hashed}:{model}"), + "model_per_team": counter("model_per_team", f"{team}:{model}"), + } + body: Final[dict[str, JsonValue]] = {} + if organization is not None: + scopes["organization"] = counter("organization", organization) + scopes["model_per_organization"] = counter("model_per_organization", f"{organization}:{model}") + if project is not None: + scopes["model_per_project"] = counter("model_per_project", f"{project}:{model}") + if cell.customer is not None: + customer: Final = f"integration-customer-{uuid.uuid4().hex}" + gateway.post("/customer/new", {"user_id": customer, **cell.customer(model)}) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [customer]}) + scopes["end_user"] = counter("end_user", customer) + body["user"] = customer + if cell.tag is not None: + tag: Final = f"integration-tag-{uuid.uuid4().hex}" + gateway.post("/tag/new", {"name": tag, **cell.tag(model)}) + scenario.cleanups.callback(gateway.request, "POST", "/tag/delete", {"name": tag}) + scopes["tag"] = counter("tag", tag) + body["metadata"] = {"tags": [tag]} + return Caller(key, user, team, scopes, body) + + +def key_rpm_only(_model: str) -> Mapping[str, JsonValue]: + return {"rpm_limit": 1000} + + +CELLS: Final = ( + Cell("organization_rpm_only", charged="organization", organization=lambda _model: {"rpm_limit": 1000}), + Cell( + "organization_model_rpm_only", + charged="model_per_organization", + organization=lambda model: {"model_rpm_limit": {model: 1000}}, + ), + Cell("key_rpm_only", charged="api_key", key=key_rpm_only), + Cell("user_rpm_only", charged="user", user=lambda _model: {"rpm_limit": 1000}), + Cell("team_rpm_only", charged="team", team=lambda _model: {"rpm_limit": 1000}), + Cell("team_member_rpm_only", charged="team_member", team=lambda _model: {"team_member_rpm_limit": 1000}), + Cell("key_model_rpm_only", charged="model_per_key", key=lambda model: {"model_rpm_limit": {model: 1000}}), + Cell( + "team_model_rpm_only", + charged="model_per_team", + also_charged=frozenset({"model_per_key"}), + team=lambda model: {"model_rpm_limit": {model: 1000}}, + ), + Cell( + "project_model_rpm_only", charged="model_per_project", project=lambda model: {"model_rpm_limit": {model: 1000}} + ), + Cell("end_user_rpm_only", charged="end_user", customer=lambda _model: {"rpm_limit": 1000}), + Cell("tag_tpm", charged="tag", tag=lambda _model: {"tpm_limit": 100000}), + Cell("key_tpm_reserved", charged="api_key", key=lambda _model: {"tpm_limit": 100000}), + Cell( + "key_dynamic_tpm_on_a_healthy_model", + charged="user", + key=lambda _model: {"tpm_limit": 100000, "tpm_limit_type": "dynamic"}, + user=lambda _model: {"rpm_limit": 1000}, + ), +) + + +@contextmanager +def redis_counters() -> Iterator[Redis]: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache: + yield cache + + +def assert_only_charged(cache: Redis, scopes: Mapping[str, str], charged: Set[str], tokens: float) -> None: + counters: Final = eventually( + lambda: read_counters(cache, scopes), + lambda values: all(values[name] == tokens for name in charged), + seconds=30, + ) + assert {name: value for name, value in counters.items() if name not in charged} == { + name: None for name in scopes if name not in charged + } + + +def chat(gateway: Gateway, identity: Caller, model: str, prompt: str, **extra: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": prompt}], **identity.body, **extra}, + key=identity.key, + ) + + +@pytest.mark.parametrize("cell", CELLS, ids=[cell.name for cell in CELLS]) +def test_success_charges_token_counters_only_for_rate_limited_scopes(gateway: Gateway, cell: Cell) -> None: + prompt: Final = f"token counters {cell.name} {uuid.uuid4().hex}" + with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, cell, model) + response: Final = chat(gateway, identity, model, prompt, max_tokens=500) + assert response.status_code == 200, response.text + assert response.json()["usage"]["total_tokens"] == USAGE["total_tokens"] + assert [json.loads(seen.body)["messages"] for seen in upstream.drain()] == [ + [{"role": "user", "content": prompt}] + ] + + assert_only_charged(cache, identity.scopes, {cell.charged, *cell.also_charged}, USAGE["total_tokens"]) + + +def gateway_url(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def openai_chat_sync(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + openai.OpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + response: Final = client.chat.completions.create(model=model, messages=[{"role": "user", "content": prompt}]) + return None if response.usage is None else response.usage.total_tokens + + +async def openai_chat_stream_async(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + openai.AsyncOpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + stream: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + stream=True, + stream_options={"include_usage": True}, + ) + usage: Final = [chunk.usage.total_tokens async for chunk in stream if chunk.usage is not None] + return usage[-1] if usage else None + + +def anthropic_messages_sync(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + anthropic.Anthropic(api_key=key, base_url=gateway_url(gateway), max_retries=0, http_client=transport) as client, + ): + message: Final = client.messages.create( + model=model, max_tokens=50, messages=[{"role": "user", "content": prompt}] + ) + return message.usage.input_tokens + message.usage.output_tokens + + +async def anthropic_messages_stream_async(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + anthropic.AsyncAnthropic( + api_key=key, base_url=gateway_url(gateway), max_retries=0, http_client=transport + ) as client, + client.messages.stream(model=model, max_tokens=50, messages=[{"role": "user", "content": prompt}]) as stream, + ): + message: Final = await stream.get_final_message() + return message.usage.input_tokens + message.usage.output_tokens + + +def openai_responses_sync(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + openai.OpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + response: Final = client.responses.create(model=model, input=prompt) + return None if response.usage is None else response.usage.total_tokens + + +async def openai_responses_stream_async(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + openai.AsyncOpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + stream: Final = await client.responses.create(model=model, input=prompt, stream=True) + usage: Final = [ + event.response.usage.total_tokens + async for event in stream + if event.type == "response.completed" and event.response.usage is not None + ] + return usage[-1] if usage else None + + +ENDPOINTS: Final[Mapping[str, Callable[[Gateway, str, str, str], int | None]]] = { + "chat_sync": openai_chat_sync, + "chat_stream_async": lambda *args: asyncio.run(openai_chat_stream_async(*args)), + "messages_sync": anthropic_messages_sync, + "messages_stream_async": lambda *args: asyncio.run(anthropic_messages_stream_async(*args)), + "responses_sync": openai_responses_sync, + "responses_stream_async": lambda *args: asyncio.run(openai_responses_stream_async(*args)), +} + + +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_every_endpoint_charges_only_the_rate_limited_scope(gateway: Gateway, endpoint: str) -> None: + prompt: Final = f"token counters {endpoint} {uuid.uuid4().hex}" + with wire_server(openai_endpoints) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, Cell(endpoint, charged="api_key", key=key_rpm_only), model) + + assert ENDPOINTS[endpoint](gateway, identity.key, model, prompt) == WIRE_TOTAL_TOKENS + assert len(upstream.drain()) == 1 + + assert_only_charged(cache, identity.scopes, {"api_key"}, WIRE_TOTAL_TOKENS) + + +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, + wire_server(completion) as fallback_upstream, + gateway.scenario() as scenario, + redis_counters() as cache, + ): + primary: Final = scenario.model(api_base=primary_upstream.url + "/v1") + fallback: Final = scenario.model(api_base=fallback_upstream.url + "/v1") + cell: Final = Cell( + "fallback", + charged="model_per_key", + 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()) == 2 + + hashed: Final = sha256(identity.key.encode()).hexdigest() + scopes: Final = { + name: identity.scopes[name] for name in ("api_key", "user", "team", "team_member", "model_per_key") + } + assert_only_charged( + cache, + {**scopes, "primary_model_per_key": counter("model_per_key", f"{hashed}:{primary}")}, + {"model_per_key"}, + 2 * USAGE["total_tokens"], + ) + + +def same_key(_scenario: Scenario, identity: Caller) -> Caller: + return identity + + +def another_key_of_the_same_user(scenario: Scenario, identity: Caller) -> Caller: + return replace(identity, key=scenario.key(user_id=identity.user, team_id=identity.team)) + + +@pytest.mark.parametrize( + "cell, next_caller", + [ + (Cell("key_tpm", charged="api_key", key=lambda _model: {"tpm_limit": 21}), same_key), + ( + Cell("user_tpm_across_keys", charged="user", user=lambda _model: {"tpm_limit": 21}), + another_key_of_the_same_user, + ), + ], + ids=["key_tpm", "user_tpm_across_keys"], +) +def test_rate_limited_scope_rejects_once_charged_tokens_reach_its_limit( + gateway: Gateway, cell: Cell, next_caller: Callable[[Scenario, Caller], Caller] +) -> None: + with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, cell, model) + first: Final = chat(gateway, identity, model, "hi", max_tokens=1) + assert first.status_code == 200, first.text + eventually( + lambda: read_counters(cache, identity.scopes)[cell.charged], + lambda tokens: tokens == USAGE["total_tokens"], + seconds=30, + ) + + rejected: Final = chat(gateway, next_caller(scenario, identity), model, "hi again", max_tokens=1) + assert rejected.status_code == 429, rejected.text + assert len(upstream.drain()) == 1 + assert_only_charged(cache, identity.scopes, {cell.charged}, USAGE["total_tokens"]) + + +def test_concurrent_burst_charges_every_request_to_the_rate_limited_scope_only(gateway: Gateway) -> None: + requests: Final = 20 + with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, Cell("burst", charged="api_key", key=key_rpm_only), model) + with ThreadPoolExecutor(max_workers=requests) as pool: + statuses: Final = list( + pool.map( + lambda index: chat(gateway, identity, model, f"token counters burst {index} {uuid.uuid4().hex}"), + range(requests), + ) + ) + assert [response.status_code for response in statuses] == [200] * requests + assert len(upstream.drain()) == requests + + assert_only_charged(cache, identity.scopes, {"api_key"}, requests * USAGE["total_tokens"]) + + +def test_upstream_outage_charges_no_token_counter(gateway: Gateway) -> None: + with wire_server(outage) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, Cell("outage", charged="api_key", key=key_rpm_only), model) + response: Final = chat(gateway, identity, model, f"token counters outage {uuid.uuid4().hex}") + assert response.status_code == 500, response.text + assert len(upstream.drain()) == 1 + + held( + lambda: read_counters(cache, identity.scopes), + 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: + with ( + wire_server(completion) as upstream, + owned_redis(tmp_path / "redis") as private_redis, + owned_proxy( + gateway, + tmp_path, + {"REDIS_HOST": private_redis.host, "REDIS_PORT": str(private_redis.port)}, + config=config_with_user_header_mapping(tmp_path), + ) as candidate, + candidate.scenario() as scenario, + Redis(host=private_redis.host, port=private_redis.port) as cache, + ): + model: Final = scenario.model(api_base=upstream.url + "/v1") + user: Final = limited_user(scenario, 100000) + user_counter: Final = {"user": counter("user", user)} + + 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 4b73bb6dd34..b568723dc56 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -26,8 +26,10 @@ from litellm.caching.redis_batch import ( from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT, PARALLEL_RELEASE_SCRIPT, TOKEN_INCREMENT_SCRIPT, + OpenWindowIncrementOperation, ParallelSlotAcquisition, PROXY_MaxParallelRequestsHandler_v3, RequestRateLimiterStash, @@ -227,6 +229,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 345f37f1db4..4a6696a3d27 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -9067,3 +9067,106 @@ async def test_a_successful_extraction_counts_its_input_tokens_toward_the_key_tp call_type="aextraction", ) assert refused.value.status_code == 429 + + +_SCOPE_TOKEN_KEYS: Final = frozenset( + { + "{api_key:key-hash}:tokens", + "{user:u}:tokens", + "{team:t}:tokens", + "{team_member:t:u}:tokens", + "{model_per_key:key-hash:gpt-4o-mini}:tokens", + "{model_per_team:t:gpt-4o-mini}:tokens", + } +) + + +def _scoped_success_kwargs(**body: object) -> dict[str, object]: + return { + "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": "gpt-4o-mini", + "user_api_key_metadata": {}, + "user_api_key_team_metadata": {}, + } + }, + "model": "gpt-4o-mini", + **body, + } + + +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": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 20}, + call_type="acompletion", + ) + + +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 {operation["key"]: operation.get("open_window_key") for operation in operations} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "limits, reserved", + [ + ({}, frozenset()), + ({"rpm_limit": 100}, frozenset()), + ({"user_rpm_limit": 100}, frozenset()), + ({"tpm_limit": 1000}, frozenset({"{api_key:key-hash}:tokens"})), + ({"team_tpm_limit": 1000}, frozenset({"{team:t}:tokens"})), + ], + ids=["no_limits", "key_rpm", "user_rpm", "key_tpm_reserved", "team_tpm_reserved"], +) +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)) + + 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 +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="123", end_user_rpm_limit=100) + ) + + assert _charged_windows(handler, _scoped_success_kwargs(user=123))["{end_user:123}:tokens"] == "{end_user:123}:window" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("parent_admitted", [True, False], ids=["sub_call_of_an_admitted_request", "no_admission"]) +async def test_success_of_a_call_that_skipped_admission_charges_every_scope_without_waiting_for_a_window( + parent_admitted: bool, +) -> None: + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + if parent_admitted: + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t"), + cache=handler.internal_usage_cache.dual_cache, + data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "litellm_call_id": "parent", + }, + call_type="acompletion", + ) + + assert _charged_windows(handler, _scoped_success_kwargs(litellm_call_id="sub-call")) == dict.fromkeys( + _SCOPE_TOKEN_KEYS + )