mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
486 lines
13 KiB
Python
486 lines
13 KiB
Python
# What is this?
|
|
## Unit tests for 'dynamic_rate_limiter.py`
|
|
import asyncio
|
|
import random
|
|
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()
|
|
|
|
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
|
|
|
|
|