mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
075781568d
Three groups, all verified by running the suite rather than by inspection. 18 files whose every test function carries an unconditional @pytest.mark.skip, 39 test functions in total. They are collected on every CI run and always skip, so they advertise coverage the suite does not have. Reasons on the marks include "AWS Suspended Account", "lakera deprecated their v1 endpoint" and "moved to using 'otel' for logging"; 26 of the marks predate 2025. 30 test functions with a byte-identical body and identical decorators to a sibling in the same file and class, differing only in name. Deleting one of each pair removes no coverage. Four further candidates were excluded because they override an inherited test, where deleting the override un-shadows the base class implementation instead of removing a duplicate. 9 test functions that a later definition of the same name shadows, so Python never binds them and pytest cannot collect them. One file that is a demo script rather than a test; its own docstring says to run it with python. Verification: collecting the 26 edited files gives 2,492 node IDs before and 2,462 after. The 30 duplicate deletions account for exactly 30 removals, the 9 shadowed deletions account for 0 (confirming at runtime that they were never collectable), nothing unexplained disappeared, and nothing new appeared. No other test or module imports any deleted symbol.
494 lines
14 KiB
Python
494 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()
|
|
import os
|
|
|
|
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
|
|
|
|
try:
|
|
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",
|
|
)
|
|
pytest.fail("Expected this to raise HTTPexception")
|
|
except HTTPException as e:
|
|
assert e.status_code == 429 # check if rate limit error raised
|
|
pass
|
|
|
|
|
|
@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
|
|
|
|
|