mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
perf(rate-limiter): write v3 token counters only into open rate limit windows
The previous revision charged tokens only to the scopes admission checked, which stopped master key, JWT, header-mapped and numeric `user` traffic from counting toward limits that other requests of the same scope enforce. Every unreserved success charge now carries its scope's window key and goes through OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT, which increments the counter only while that window exists. A scope with no limit never opens a window, so it gets no counter; a counter written outside a window was deleted by the next admission anyway, so every limit sees the same token totals as before. A hash-tag group that mixes plain and gated charges goes out as one script call, so the group lands or falls back as a whole. Reserved scopes keep their reconciliation deltas, and the no-Redis path still increments in memory as before. Resolves LIT-9875
This commit is contained in:
yucheng-berri
parent
8b631a31e0
commit
8f9ba9b987
5 changed files with 346 additions and 191 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue