mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
243ed4393d
A test that asserts on the error inside its own except block passes when the call stops raising, because nothing runs the handler. That is the exact case the test exists to catch, so the regression lands green. Rewrites all 111 such blocks into pytest.raises, which fails when the call succeeds, and selects PT017 in ruff-tests.toml so no new one lands.
491 lines
14 KiB
Python
491 lines
14 KiB
Python
# What is this?
|
|
## Unit tests for 'dynamic_rate_limiter.py`
|
|
import asyncio
|
|
import os
|
|
import random
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from litellm._uuid import uuid
|
|
from datetime import datetime, timezone
|
|
from typing import Optional, Tuple
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm import DualCache, Router
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.hooks.dynamic_rate_limiter import (
|
|
_PROXY_DynamicRateLimitHandler as DynamicRateLimitHandler,
|
|
)
|
|
|
|
"""
|
|
Basic test cases:
|
|
|
|
- If 1 'active' project => give all tpm
|
|
- If 2 'active' projects => divide tpm in 2
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def dynamic_rate_limit_handler() -> DynamicRateLimitHandler:
|
|
internal_cache = DualCache()
|
|
frozen_now = datetime(2024, 1, 1, 10, 30, 0, tzinfo=timezone.utc)
|
|
return DynamicRateLimitHandler(internal_usage_cache=internal_cache, time_fn=lambda: frozen_now)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_response() -> litellm.ModelResponse:
|
|
return litellm.ModelResponse(
|
|
**{
|
|
"id": "chatcmpl-abc123",
|
|
"object": "chat.completion",
|
|
"created": 1699896916,
|
|
"model": "gpt-3.5-turbo-0125",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_abc123",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_current_weather",
|
|
"arguments": '{\n"location": "Boston, MA"\n}',
|
|
},
|
|
}
|
|
],
|
|
},
|
|
"logprobs": None,
|
|
"finish_reason": "tool_calls",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 5, "completion_tokens": 5, "total_tokens": 10},
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def user_api_key_auth() -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth()
|
|
|
|
|
|
@pytest.mark.parametrize("num_projects", [1, 2, 100])
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.flaky(retries=3, delay=1)
|
|
async def test_available_tpm(num_projects, dynamic_rate_limit_handler):
|
|
model = "my-fake-model"
|
|
## SET CACHE W/ ACTIVE PROJECTS
|
|
projects = [str(uuid.uuid4()) for _ in range(num_projects)]
|
|
|
|
await dynamic_rate_limit_handler.internal_usage_cache.async_set_cache_sadd(
|
|
model=model, value=projects
|
|
)
|
|
|
|
model_tpm = 100
|
|
llm_router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": model,
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": "my-key",
|
|
"api_base": "my-base",
|
|
"tpm": model_tpm,
|
|
},
|
|
}
|
|
]
|
|
)
|
|
dynamic_rate_limit_handler.update_variables(llm_router=llm_router)
|
|
|
|
## CHECK AVAILABLE TPM PER PROJECT
|
|
|
|
resp = await dynamic_rate_limit_handler.check_available_usage(model=model)
|
|
|
|
availability = resp[0]
|
|
|
|
expected_availability = int(model_tpm / num_projects)
|
|
|
|
assert availability == expected_availability
|
|
|
|
|
|
@pytest.mark.parametrize("num_projects", [1, 2, 100])
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.flaky(retries=3, delay=1)
|
|
async def test_available_rpm(num_projects, dynamic_rate_limit_handler):
|
|
model = "my-fake-model"
|
|
## SET CACHE W/ ACTIVE PROJECTS
|
|
projects = [str(uuid.uuid4()) for _ in range(num_projects)]
|
|
|
|
await dynamic_rate_limit_handler.internal_usage_cache.async_set_cache_sadd(
|
|
model=model, value=projects
|
|
)
|
|
|
|
model_rpm = 100
|
|
llm_router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": model,
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": "my-key",
|
|
"api_base": "my-base",
|
|
"rpm": model_rpm,
|
|
},
|
|
}
|
|
]
|
|
)
|
|
dynamic_rate_limit_handler.update_variables(llm_router=llm_router)
|
|
|
|
## CHECK AVAILABLE rpm PER PROJECT
|
|
|
|
resp = await dynamic_rate_limit_handler.check_available_usage(model=model)
|
|
|
|
availability = resp[1]
|
|
|
|
expected_availability = int(model_rpm / num_projects)
|
|
|
|
assert availability == expected_availability
|
|
|
|
|
|
@pytest.mark.parametrize("usage", ["rpm", "tpm"])
|
|
@pytest.mark.asyncio
|
|
async def test_rate_limit_raised(dynamic_rate_limit_handler, user_api_key_auth, usage):
|
|
"""
|
|
Unit test. Tests if rate limit error raised when quota exhausted.
|
|
"""
|
|
from fastapi import HTTPException
|
|
|
|
model = "my-fake-model"
|
|
## SET CACHE W/ ACTIVE PROJECTS
|
|
projects = [str(uuid.uuid4())]
|
|
|
|
await dynamic_rate_limit_handler.internal_usage_cache.async_set_cache_sadd(
|
|
model=model, value=projects
|
|
)
|
|
|
|
model_usage = 0
|
|
llm_router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": model,
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": "my-key",
|
|
"api_base": "my-base",
|
|
usage: model_usage,
|
|
},
|
|
}
|
|
]
|
|
)
|
|
dynamic_rate_limit_handler.update_variables(llm_router=llm_router)
|
|
|
|
## CHECK AVAILABLE TPM PER PROJECT
|
|
|
|
resp = await dynamic_rate_limit_handler.check_available_usage(model=model)
|
|
|
|
if usage == "tpm":
|
|
availability = resp[0]
|
|
else:
|
|
availability = resp[1]
|
|
|
|
expected_availability = 0
|
|
|
|
assert availability == expected_availability
|
|
|
|
## CHECK if exception raised
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await dynamic_rate_limit_handler.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_auth,
|
|
cache=DualCache(),
|
|
data={"model": model},
|
|
call_type="completion",
|
|
)
|
|
e = exc_info.value
|
|
assert e.status_code == 429 # check if rate limit error raised
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_base_case(dynamic_rate_limit_handler, mock_response):
|
|
"""
|
|
If just 1 active project
|
|
|
|
it should get all the quota
|
|
|
|
= allow request to go through
|
|
- update token usage
|
|
- exhaust all tpm with just 1 project
|
|
- assert ratelimiterror raised at 100%+1 tpm
|
|
"""
|
|
model = "my-fake-model"
|
|
## model tpm - 50
|
|
model_tpm = 50
|
|
## tpm per request - 10
|
|
setattr(
|
|
mock_response,
|
|
"usage",
|
|
litellm.Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10),
|
|
)
|
|
|
|
llm_router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": model,
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": "my-key",
|
|
"api_base": "my-base",
|
|
"tpm": model_tpm,
|
|
"mock_response": mock_response,
|
|
},
|
|
}
|
|
]
|
|
)
|
|
dynamic_rate_limit_handler.update_variables(llm_router=llm_router)
|
|
|
|
prev_availability: Optional[int] = None
|
|
allowed_fails = 1
|
|
for _ in range(2):
|
|
try:
|
|
# check availability
|
|
resp = await dynamic_rate_limit_handler.check_available_usage(model=model)
|
|
|
|
availability = resp[0]
|
|
|
|
print(
|
|
"prev_availability={}, availability={}".format(
|
|
prev_availability, availability
|
|
)
|
|
)
|
|
|
|
## assert availability updated
|
|
if prev_availability is not None and availability is not None:
|
|
assert availability == prev_availability - 10
|
|
|
|
prev_availability = availability
|
|
|
|
# make call
|
|
await llm_router.acompletion(
|
|
model=model, messages=[{"role": "user", "content": "hey!"}]
|
|
)
|
|
|
|
await asyncio.sleep(3)
|
|
except Exception:
|
|
if allowed_fails > 0:
|
|
allowed_fails -= 1
|
|
else:
|
|
raise
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.flaky(retries=3, delay=1)
|
|
async def test_update_cache(
|
|
dynamic_rate_limit_handler, mock_response, user_api_key_auth
|
|
):
|
|
"""
|
|
Check if active project correctly updated
|
|
"""
|
|
model = "my-fake-model"
|
|
model_tpm = 50
|
|
|
|
llm_router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": model,
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": "my-key",
|
|
"api_base": "my-base",
|
|
"tpm": model_tpm,
|
|
"mock_response": mock_response,
|
|
},
|
|
}
|
|
]
|
|
)
|
|
dynamic_rate_limit_handler.update_variables(llm_router=llm_router)
|
|
|
|
## INITIAL ACTIVE PROJECTS - ASSERT NONE
|
|
resp = await dynamic_rate_limit_handler.check_available_usage(model=model)
|
|
|
|
active_projects = resp[-1]
|
|
|
|
assert active_projects is None
|
|
|
|
## MAKE CALL
|
|
await dynamic_rate_limit_handler.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_auth,
|
|
cache=DualCache(),
|
|
data={"model": model},
|
|
call_type="completion",
|
|
)
|
|
|
|
await asyncio.sleep(2)
|
|
## INITIAL ACTIVE PROJECTS - ASSERT 1
|
|
resp = await dynamic_rate_limit_handler.check_available_usage(model=model)
|
|
|
|
active_projects = resp[-1]
|
|
|
|
assert active_projects == 1
|
|
|
|
|
|
@pytest.mark.skip(
|
|
reason="Unstable on ci/cd due to curr minute changes. Refactor to handle minute changing"
|
|
)
|
|
@pytest.mark.parametrize("num_projects", [2])
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_projects(
|
|
dynamic_rate_limit_handler, mock_response, num_projects
|
|
):
|
|
"""
|
|
If 2 active project
|
|
|
|
it should split 50% each
|
|
|
|
- assert available tpm is 0 after 50%+1 tpm calls
|
|
"""
|
|
model = "my-fake-model"
|
|
model_tpm = 50
|
|
total_tokens_per_call = 10
|
|
step_tokens_per_call_per_project = total_tokens_per_call / num_projects
|
|
|
|
available_tpm_per_project = int(model_tpm / num_projects)
|
|
|
|
## SET CACHE W/ ACTIVE PROJECTS
|
|
projects = [str(uuid.uuid4()) for _ in range(num_projects)]
|
|
await dynamic_rate_limit_handler.internal_usage_cache.async_set_cache_sadd(
|
|
model=model, value=projects
|
|
)
|
|
|
|
expected_runs = int(available_tpm_per_project / step_tokens_per_call_per_project)
|
|
|
|
setattr(
|
|
mock_response,
|
|
"usage",
|
|
litellm.Usage(
|
|
prompt_tokens=5, completion_tokens=5, total_tokens=total_tokens_per_call
|
|
),
|
|
)
|
|
|
|
llm_router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": model,
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": "my-key",
|
|
"api_base": "my-base",
|
|
"tpm": model_tpm,
|
|
"mock_response": mock_response,
|
|
},
|
|
}
|
|
]
|
|
)
|
|
dynamic_rate_limit_handler.update_variables(llm_router=llm_router)
|
|
|
|
prev_availability: Optional[int] = None
|
|
|
|
print("expected_runs: {}".format(expected_runs))
|
|
|
|
for i in range(expected_runs + 1):
|
|
# check availability
|
|
|
|
resp = await dynamic_rate_limit_handler.check_available_usage(model=model)
|
|
|
|
availability = resp[0]
|
|
|
|
## assert availability updated
|
|
if prev_availability is not None and availability is not None:
|
|
assert (
|
|
availability == prev_availability - step_tokens_per_call_per_project
|
|
), "Current Availability: Got={}, Expected={}, Step={}, Tokens per step={}, Initial model tpm={}".format(
|
|
availability,
|
|
prev_availability - 10,
|
|
i,
|
|
step_tokens_per_call_per_project,
|
|
model_tpm,
|
|
)
|
|
|
|
print(
|
|
"prev_availability={}, availability={}".format(
|
|
prev_availability, availability
|
|
)
|
|
)
|
|
|
|
prev_availability = availability
|
|
|
|
# make call
|
|
await llm_router.acompletion(
|
|
model=model, messages=[{"role": "user", "content": "hey!"}]
|
|
)
|
|
|
|
await asyncio.sleep(3)
|
|
|
|
# check availability
|
|
resp = await dynamic_rate_limit_handler.check_available_usage(model=model)
|
|
|
|
availability = resp[0]
|
|
|
|
assert availability == 0
|
|
|
|
|
|
@pytest.mark.parametrize("num_projects", [1, 2, 100])
|
|
@pytest.mark.asyncio
|
|
async def test_priority_reservation(num_projects, dynamic_rate_limit_handler):
|
|
"""
|
|
If reservation is set + `mock_testing_reservation` passed in
|
|
|
|
assert correct rpm is reserved
|
|
"""
|
|
model = "my-fake-model"
|
|
## SET CACHE W/ ACTIVE PROJECTS
|
|
projects = [str(uuid.uuid4()) for _ in range(num_projects)]
|
|
|
|
await dynamic_rate_limit_handler.internal_usage_cache.async_set_cache_sadd(
|
|
model=model, value=projects
|
|
)
|
|
|
|
litellm.priority_reservation = {"dev": 0.1, "prod": 0.9}
|
|
|
|
model_usage = 100
|
|
|
|
llm_router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": model,
|
|
"litellm_params": {
|
|
"model": "gpt-3.5-turbo",
|
|
"api_key": "my-key",
|
|
"api_base": "my-base",
|
|
"rpm": model_usage,
|
|
},
|
|
}
|
|
]
|
|
)
|
|
dynamic_rate_limit_handler.update_variables(llm_router=llm_router)
|
|
|
|
## CHECK AVAILABLE TPM PER PROJECT
|
|
|
|
resp = await dynamic_rate_limit_handler.check_available_usage(
|
|
model=model, priority="prod"
|
|
)
|
|
|
|
availability = resp[1]
|
|
|
|
expected_availability = int(
|
|
model_usage * litellm.priority_reservation["prod"] / num_projects
|
|
)
|
|
|
|
assert availability == expected_availability
|
|
|
|
|