scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87
litellm/tests/local_testing/test_lowest_latency_routing.py
yuneng-jiang 6a0d03914c
test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 
* 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
2026-08-22 09:25:58 -07:00

1340 lines
41 KiB
Python
Raw Permalink Blame History

#### What this tests ####
# This tests the router's ability to pick deployment with lowest latency
import asyncio
import random
import sys
import time
import traceback
from datetime import datetime, timedelta
from dotenv import load_dotenv
load_dotenv()
import copy
import pytest
import litellm
from litellm import Router
from litellm.caching.caching import DualCache
from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
### UNIT TESTS FOR LATENCY ROUTING ###
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_latency_memory_leak(sync_mode):
"""
Test to make sure there's no memory leak caused by lowest latency routing
- make 10 calls -> check memory
- make 11th call -> no change in memory
"""
test_cache = DualCache()
lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
model_group = "gpt-3.5-turbo"
deployment_id = "1234"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-3.5-turbo",
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
start_time = time.time()
response_obj = {"usage": {"total_tokens": 50}}
time.sleep(5)
end_time = time.time()
for _ in range(10):
if sync_mode:
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
else:
await lowest_latency_logger.async_log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_map"
cache_value = copy.deepcopy(
test_cache.get_cache(key=latency_key)
) # MAKE SURE NO MEMORY LEAK IN CACHING OBJECT
if sync_mode:
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
else:
await lowest_latency_logger.async_log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
new_cache_value = test_cache.get_cache(key=latency_key)
# Assert that the size of the cache doesn't grow unreasonably
assert get_size(new_cache_value) <= get_size(
cache_value
), f"Memory leak detected in function call! new_cache size={get_size(new_cache_value)}, old cache size={get_size(cache_value)}"
def get_size(obj, seen=None):
# From https://goshippo.com/blog/measure-real-size-any-python-object/
# Recursively finds size of objects
size = sys.getsizeof(obj)
if seen is None:
seen = set()
obj_id = id(obj)
if obj_id in seen:
return 0
seen.add(obj_id)
if isinstance(obj, dict):
size += sum([get_size(v, seen) for v in obj.values()])
size += sum([get_size(k, seen) for k in obj.keys()])
elif hasattr(obj, "__dict__"):
size += get_size(obj.__dict__, seen)
elif hasattr(obj, "__iter__") and not isinstance(obj, (str, bytes, bytearray)):
size += sum([get_size(i, seen) for i in obj])
return size
def test_latency_updated():
test_cache = DualCache()
lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
model_group = "gpt-3.5-turbo"
deployment_id = "1234"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-3.5-turbo",
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
start_time = time.time()
response_obj = {"usage": {"total_tokens": 50}}
time.sleep(5)
end_time = time.time()
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_map"
assert (
end_time - start_time
== test_cache.get_cache(key=latency_key)[deployment_id]["latency"][0]
)
# test_tpm_rpm_updated()
def test_latency_updated_custom_ttl():
"""
Invalidate the cached request.
Test that the cache is empty
"""
test_cache = DualCache()
model_list = []
cache_time = 3
lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, routing_args={"ttl": cache_time}
)
model_group = "gpt-3.5-turbo"
deployment_id = "1234"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-3.5-turbo",
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
start_time = time.time()
response_obj = {"usage": {"total_tokens": 50}}
time.sleep(5)
end_time = time.time()
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_map"
print(f"cache: {test_cache.get_cache(key=latency_key)}")
assert isinstance(test_cache.get_cache(key=latency_key), dict)
time.sleep(cache_time)
assert test_cache.get_cache(key=latency_key) is None
def test_get_available_deployments():
test_cache = DualCache()
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "azure/gpt-4.1-mini"},
"model_info": {"id": "1234"},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "azure/gpt-4.1-mini"},
"model_info": {"id": "5678"},
},
]
lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
model_group = "gpt-3.5-turbo"
## DEPLOYMENT 1 ##
deployment_id = "1234"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-3.5-turbo",
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
start_time = time.time()
response_obj = {"usage": {"total_tokens": 50}}
time.sleep(3)
end_time = time.time()
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
## DEPLOYMENT 2 ##
deployment_id = "5678"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-3.5-turbo",
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
start_time = time.time()
response_obj = {"usage": {"total_tokens": 20}}
time.sleep(2)
end_time = time.time()
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
## CHECK WHAT'S SELECTED ##
print(
lowest_latency_logger.get_available_deployments(
model_group=model_group, healthy_deployments=model_list
)
)
assert (
lowest_latency_logger.get_available_deployments(
model_group=model_group, healthy_deployments=model_list
)["model_info"]["id"]
== "5678"
)
async def _deploy(lowest_latency_logger, deployment_id, tokens_used, duration):
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-3.5-turbo",
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
start_time = time.time()
response_obj = {"usage": {"total_tokens": tokens_used}}
await asyncio.sleep(duration)
end_time = time.time()
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
async def _gather_deploy(all_deploys):
return await asyncio.gather(*[_deploy(*t) for t in all_deploys])
@pytest.mark.parametrize(
"ans_rpm", [1, 5]
) # 1 should produce nothing, 10 should select first
@pytest.mark.flaky(retries=3, delay=1)
def test_get_available_endpoints_tpm_rpm_check_async(ans_rpm):
"""
Pass in list of 2 valid models
Update cache with 1 model clearly being at tpm/rpm limit
assert that only the valid model is returned
"""
test_cache = DualCache()
ans = "1234"
non_ans_rpm = 3
assert ans_rpm != non_ans_rpm, "invalid test"
if ans_rpm < non_ans_rpm:
ans = None
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "azure/gpt-4.1-mini"},
"model_info": {"id": "1234", "rpm": ans_rpm},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "azure/gpt-4.1-mini"},
"model_info": {"id": "5678", "rpm": non_ans_rpm},
},
]
lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
model_group = "gpt-3.5-turbo"
d1 = [(lowest_latency_logger, "1234", 50, 0.01)] * non_ans_rpm
d2 = [(lowest_latency_logger, "5678", 50, 0.01)] * non_ans_rpm
asyncio.run(_gather_deploy([*d1, *d2]))
time.sleep(3)
## CHECK WHAT'S SELECTED ##
d_ans = lowest_latency_logger.get_available_deployments(
model_group=model_group, healthy_deployments=model_list
)
print(d_ans)
assert (d_ans and d_ans["model_info"]["id"]) == ans
# test_get_available_endpoints_tpm_rpm_check_async()
@pytest.mark.parametrize(
"ans_rpm", [1, 5]
) # 1 should produce nothing, 10 should select first
@pytest.mark.flaky(retries=3, delay=1)
def test_get_available_endpoints_tpm_rpm_check(ans_rpm):
"""
Pass in list of 2 valid models
Update cache with 1 model clearly being at tpm/rpm limit
assert that only the valid model is returned
"""
test_cache = DualCache()
ans = "1234"
non_ans_rpm = 3
assert ans_rpm != non_ans_rpm, "invalid test"
if ans_rpm < non_ans_rpm:
ans = None
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "azure/gpt-4.1-mini"},
"model_info": {"id": "1234", "rpm": ans_rpm},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "azure/gpt-4.1-mini"},
"model_info": {"id": "5678", "rpm": non_ans_rpm},
},
]
lowest_latency_logger = LowestLatencyLoggingHandler(router_cache=test_cache)
model_group = "gpt-3.5-turbo"
## DEPLOYMENT 1 ##
deployment_id = "1234"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-3.5-turbo",
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
for _ in range(non_ans_rpm):
start_time = time.time()
response_obj = {"usage": {"total_tokens": 50}}
time.sleep(0.01)
end_time = time.time()
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
## DEPLOYMENT 2 ##
deployment_id = "5678"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-3.5-turbo",
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
for _ in range(non_ans_rpm):
start_time = time.time()
response_obj = {"usage": {"total_tokens": 20}}
time.sleep(0.5)
end_time = time.time()
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
## CHECK WHAT'S SELECTED ##
d_ans = lowest_latency_logger.get_available_deployments(
model_group=model_group, healthy_deployments=model_list
)
print(d_ans)
assert (d_ans and d_ans["model_info"]["id"]) == ans
def test_router_get_available_deployments():
"""
Test if routers 'get_available_deployments' returns the fastest deployment
"""
model_list = [
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-turbo",
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
"api_base": "https://openai-france-1234.openai.azure.com",
"rpm": 1440,
},
"model_info": {"id": 1},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
"rpm": 6,
},
"model_info": {"id": 2},
},
]
router = Router(
model_list=model_list,
routing_strategy="latency-based-routing",
set_verbose=False,
num_retries=3,
) # type: ignore
## DEPLOYMENT 1 ##
deployment_id = 1
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "azure-model",
},
"model_info": {"id": 1},
}
}
start_time = time.time()
response_obj = {"usage": {"total_tokens": 50}}
time.sleep(3)
end_time = time.time()
router.lowestlatency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
## DEPLOYMENT 2 ##
deployment_id = 2
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "azure-model",
},
"model_info": {"id": 2},
}
}
start_time = time.time()
response_obj = {"usage": {"total_tokens": 20}}
time.sleep(2)
end_time = time.time()
router.lowestlatency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
## CHECK WHAT'S SELECTED ##
# print(router.lowesttpm_logger.get_available_deployments(model_group="azure-model"))
print(router.get_available_deployment(model="azure-model"))
assert (
router.get_available_deployment(model="azure-model")["model_info"]["id"] == "2"
)
# test_router_get_available_deployments()
@pytest.mark.asyncio
async def test_router_completion_streaming():
messages = [
{"role": "user", "content": "Hello, can you generate a 500 words poem?"}
]
model = "azure-model"
model_list = [
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-turbo",
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
"api_base": "https://openai-france-1234.openai.azure.com",
"rpm": 1440,
"mock_response": "Hello world",
},
"model_info": {"id": 1},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
"rpm": 6,
"mock_response": "Hello world",
},
"model_info": {"id": 2},
},
]
router = Router(
model_list=model_list,
routing_strategy="latency-based-routing",
set_verbose=False,
num_retries=3,
) # type: ignore
### Make 3 calls, test if 3rd call goes to fastest deployment
## CALL 1+2
tasks = []
response = None
final_response = None
for _ in range(2):
tasks.append(router.acompletion(model=model, messages=messages))
response = await asyncio.gather(*tasks)
if response is not None:
## CALL 3
await asyncio.sleep(1) # let the cache update happen
picked_deployment = router.lowestlatency_logger.get_available_deployments(
model_group=model, healthy_deployments=router.healthy_deployments
)
final_response = await router.acompletion(model=model, messages=messages)
print(f"min deployment id: {picked_deployment}")
print(f"model id: {final_response._hidden_params['model_id']}")
assert (
final_response._hidden_params["model_id"]
== picked_deployment["model_info"]["id"]
)
# asyncio.run(test_router_completion_streaming())
@pytest.mark.asyncio
async def test_lowest_latency_routing_with_timeouts():
"""
PROD Test:
- Endpoint 1: triggers timeout errors (it takes 10+ seconds to respond)
- Endpoint 2: Responds in under 1s
- Run 5 requests to collect data on latency
- Run Wait till cache is filled with data
- Run 10 more requests
- All requests should have been routed to endpoint 2
"""
import litellm
litellm.set_verbose = True
router = Router(
model_list=[
{
"model_name": "azure-model",
"litellm_params": {
"model": "openai/slow-endpoint",
"api_base": FAKE_OPENAI_API_BASE,
"api_key": "fake-key",
},
"model_info": {"id": "slow-endpoint"},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "openai/fast-endpoint",
"api_base": FAKE_OPENAI_API_BASE,
"api_key": "fake-key",
},
"model_info": {"id": "fast-endpoint"},
},
],
routing_strategy="latency-based-routing",
set_verbose=True,
debug_level="DEBUG",
timeout=1,
) # type: ignore
# make 4 requests
for _ in range(4):
try:
response = await router.acompletion(
model="azure-model", messages=[{"role": "user", "content": "hello"}]
)
print(response)
except Exception as e:
print("got exception", e)
await asyncio.sleep(1)
print("done sending initial requests to collect latency")
"""
Note: for debugging
- By this point: slow-endpoint should have timed out 3-4 times and should be heavily penalized :)
- The next 10 requests should all be routed to the fast-endpoint
"""
deployments = {}
# make 10 requests
for _ in range(10):
response = await router.acompletion(
model="azure-model", messages=[{"role": "user", "content": "hello"}]
)
print(response)
_picked_model_id = response._hidden_params["model_id"]
if _picked_model_id not in deployments:
deployments[_picked_model_id] = 1
else:
deployments[_picked_model_id] += 1
print("deployments", deployments)
# ALL the Requests should have been routed to the fast-endpoint
assert deployments["fast-endpoint"] == 10
@pytest.mark.asyncio
async def test_lowest_latency_routing_first_pick():
"""
PROD Test:
- When all deployments are latency=0, it should randomly pick a deployment
- IT SHOULD NEVER PICK THE Very First deployment everytime all deployment latencies are 0
- This ensures that after the ttl window resets it randomly picks a deployment
"""
import litellm
litellm.set_verbose = True
router = Router(
model_list=[
{
"model_name": "azure-model",
"litellm_params": {
"model": "openai/fast-endpoint",
"api_base": FAKE_OPENAI_API_BASE,
"api_key": "fake-key",
},
"model_info": {"id": "fast-endpoint"},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "openai/fast-endpoint-2",
"api_base": FAKE_OPENAI_API_BASE,
"api_key": "fake-key",
},
"model_info": {"id": "fast-endpoint-2"},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "openai/fast-endpoint-2",
"api_base": FAKE_OPENAI_API_BASE,
"api_key": "fake-key",
},
"model_info": {"id": "fast-endpoint-3"},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "openai/fast-endpoint-2",
"api_base": FAKE_OPENAI_API_BASE,
"api_key": "fake-key",
},
"model_info": {"id": "fast-endpoint-4"},
},
],
routing_strategy="latency-based-routing",
routing_strategy_args={"ttl": 0.0000000001},
set_verbose=True,
debug_level="DEBUG",
) # type: ignore
deployments = {}
for _ in range(10):
response = await router.acompletion(
model="azure-model", messages=[{"role": "user", "content": "hello"}]
)
print(response)
_picked_model_id = response._hidden_params["model_id"]
if _picked_model_id not in deployments:
deployments[_picked_model_id] = 1
else:
deployments[_picked_model_id] += 1
await asyncio.sleep(0.000000000005)
print("deployments", deployments)
# assert that len(deployments) >1
assert len(deployments) > 1
@pytest.mark.parametrize("buffer", [0, 1])
@pytest.mark.asyncio
async def test_lowest_latency_routing_buffer(buffer):
"""
Allow shuffling calls within a certain latency buffer
"""
model_list = [
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-turbo",
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
"api_base": "https://openai-france-1234.openai.azure.com",
"rpm": 1440,
},
"model_info": {"id": 1},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
"rpm": 6,
},
"model_info": {"id": 2},
},
]
router = Router(
model_list=model_list,
routing_strategy="latency-based-routing",
set_verbose=False,
num_retries=3,
routing_strategy_args={"lowest_latency_buffer": buffer},
) # type: ignore
## DEPLOYMENT 1 ##
deployment_id = 1
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "azure-model",
},
"model_info": {"id": 1},
}
}
start_time = time.time()
response_obj = {"usage": {"total_tokens": 50}}
time.sleep(3)
end_time = time.time()
router.lowestlatency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
## DEPLOYMENT 2 ##
deployment_id = 2
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "azure-model",
},
"model_info": {"id": 2},
}
}
start_time = time.time()
response_obj = {"usage": {"total_tokens": 20}}
time.sleep(2)
end_time = time.time()
router.lowestlatency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
## CHECK WHAT'S SELECTED ##
# print(router.lowesttpm_logger.get_available_deployments(model_group="azure-model"))
selected_deployments = {}
for _ in range(50):
print(router.get_available_deployment(model="azure-model"))
selected_deployments[
router.get_available_deployment(model="azure-model")["model_info"]["id"]
] = 1
if buffer == 0:
assert len(selected_deployments.keys()) == 1
else:
assert len(selected_deployments.keys()) == 2
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_lowest_latency_routing_time_to_first_token(sync_mode):
"""
If a deployment has
- a fast time to first token
- slow latency/output token
test if:
- for streaming, the deployment with fastest time to first token is picked
- for non-streaming, fastest overall deployment is picked
"""
model_list = [
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-turbo",
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
"api_base": "https://openai-france-1234.openai.azure.com",
},
"model_info": {"id": 1},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
},
"model_info": {"id": 2},
},
]
router = Router(
model_list=model_list,
routing_strategy="latency-based-routing",
set_verbose=False,
num_retries=3,
) # type: ignore
## DEPLOYMENT 1 ##
deployment_id = 1
start_time = datetime.now()
one_second_later = start_time + timedelta(seconds=1)
# Compute 3 seconds after the current time
three_seconds_later = start_time + timedelta(seconds=3)
four_seconds_later = start_time + timedelta(seconds=4)
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "azure-model",
},
"model_info": {"id": 1},
},
"stream": True,
"completion_start_time": one_second_later,
}
response_obj = litellm.ModelResponse(
usage=litellm.Usage(completion_tokens=50, total_tokens=50)
)
end_time = four_seconds_later
if sync_mode:
router.lowestlatency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
else:
await router.lowestlatency_logger.async_log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
## DEPLOYMENT 2 ##
deployment_id = 2
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "azure-model",
},
"model_info": {"id": 2},
},
"stream": True,
"completion_start_time": three_seconds_later,
}
response_obj = litellm.ModelResponse(
usage=litellm.Usage(completion_tokens=50, total_tokens=50)
)
end_time = three_seconds_later
if sync_mode:
router.lowestlatency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
else:
await router.lowestlatency_logger.async_log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
"""
TESTING
- expect deployment 1 to be picked for streaming
- expect deployment 2 to be picked for non-streaming
"""
# print(router.lowesttpm_logger.get_available_deployments(model_group="azure-model"))
selected_deployments = {}
for _ in range(3):
print(router.get_available_deployment(model="azure-model"))
## for non-streaming
selected_deployments[
router.get_available_deployment(model="azure-model")["model_info"]["id"]
] = 1
assert len(selected_deployments.keys()) == 1
assert "2" in list(selected_deployments.keys())
selected_deployments = {}
for _ in range(50):
print(router.get_available_deployment(model="azure-model"))
## for non-streaming
selected_deployments[
router.get_available_deployment(
model="azure-model", request_kwargs={"stream": True}
)["model_info"]["id"]
] = 1
assert len(selected_deployments.keys()) == 1
assert "1" in list(selected_deployments.keys())
def test_latency_list_trimming_discards_oldest_entry():
"""
When the latency list reaches max_latency_list_size, the oldest entry is
discarded to make room for new entries. The newest entry is appended at
the end of the list.
"""
max_size = 3
test_cache = DualCache()
lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
)
model_group = "gpt-3.5-turbo"
deployment_id = "test-deployment"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": model_group,
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
# With 1 completion token, the logged latency value equals the raw
# response time, so we can use distinct, identifiable values.
latencies_to_add = []
for i in range(max_size + 1): # One more than max to trigger trimming
start_time = time.time()
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
expected_latency = float(i + 1) # 1.0, 2.0, 3.0, 4.0
end_time = start_time + expected_latency
latencies_to_add.append(expected_latency)
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_map"
cached_data = test_cache.get_cache(key=latency_key)
latency_list = cached_data[deployment_id]["latency"]
assert (
len(latency_list) == max_size
), f"Expected {max_size} entries, got {len(latency_list)}"
newest_latency = latencies_to_add[-1] # 4.0
oldest_latency = latencies_to_add[0] # 1.0
tolerance = 0.1
# Newest entry is at the end of the list.
assert (
abs(latency_list[-1] - newest_latency) < tolerance
), f"Newest latency {newest_latency} should be at end, got {latency_list[-1]}"
# Oldest entry is no longer in the list.
for latency in latency_list:
assert (
abs(latency - oldest_latency) > tolerance
), f"Oldest latency {oldest_latency} should have been discarded, found {latency}"
@pytest.mark.asyncio
async def test_latency_list_trimming_discards_oldest_entry_async():
"""
Async counterpart: the oldest entry is discarded when the latency list is
trimmed.
"""
max_size = 3
test_cache = DualCache()
lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
)
model_group = "gpt-3.5-turbo"
deployment_id = "test-deployment"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": model_group,
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
latencies_to_add = []
for i in range(max_size + 1):
start_time = time.time()
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
expected_latency = float(i + 1)
end_time = start_time + expected_latency
latencies_to_add.append(expected_latency)
await lowest_latency_logger.async_log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_map"
cached_data = await test_cache.async_get_cache(key=latency_key)
latency_list = cached_data[deployment_id]["latency"]
assert len(latency_list) == max_size
newest_latency = latencies_to_add[-1]
oldest_latency = latencies_to_add[0]
tolerance = 0.1
assert (
abs(latency_list[-1] - newest_latency) < tolerance
), f"Newest latency {newest_latency} should be at end of list"
for latency in latency_list:
assert (
abs(latency - oldest_latency) > tolerance
), f"Oldest latency {oldest_latency} should have been discarded"
def test_ttft_list_trimming_discards_oldest_entry():
"""
The time_to_first_token list trims the oldest entry when full, matching
the behavior of the latency list.
"""
max_size = 3
test_cache = DualCache()
lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
)
model_group = "gpt-3.5-turbo"
deployment_id = "test-deployment"
ttft_values = []
for i in range(max_size + 1):
start_time = time.time()
expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4
completion_start_time = start_time + expected_ttft
end_time = start_time + float(i + 1)
ttft_values.append(expected_ttft)
kwargs = {
"litellm_params": {
"metadata": {
"model_group": model_group,
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
},
"stream": True,
"completion_start_time": completion_start_time,
}
# TTFT is only recorded when response_obj is a ModelResponse.
response_obj = litellm.ModelResponse(
usage=litellm.Usage(completion_tokens=1, total_tokens=1)
)
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_map"
cached_data = test_cache.get_cache(key=latency_key)
ttft_list = cached_data[deployment_id].get("time_to_first_token", [])
assert (
len(ttft_list) == max_size
), f"Expected {max_size} entries, got {len(ttft_list)}"
newest_ttft = ttft_values[-1]
oldest_ttft = ttft_values[0]
tolerance = 0.05
assert (
abs(ttft_list[-1] - newest_ttft) < tolerance
), f"Newest TTFT {newest_ttft} should be at end of list"
for ttft in ttft_list:
assert (
abs(ttft - oldest_ttft) > tolerance
), f"Oldest TTFT {oldest_ttft} should have been discarded"
@pytest.mark.asyncio
async def test_timeout_penalty_discards_oldest_entry():
"""
Timeout penalties (1000.0) are appended to the latency list and, when the
list is full, the oldest entry is discarded.
"""
max_size = 3
test_cache = DualCache()
lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
)
model_group = "gpt-3.5-turbo"
deployment_id = "test-deployment"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": model_group,
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
# Fill the list with max_size normal latency entries first.
for i in range(max_size):
start_time = time.time()
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
end_time = start_time + float(i + 1)
await lowest_latency_logger.async_log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
# Trigger a timeout failure: this appends 1000.0 and should discard the
# oldest normal entry (1.0).
timeout_kwargs = {
**kwargs,
"exception": litellm.Timeout(
message="Request timed out", model="test-model", llm_provider="test"
),
}
await lowest_latency_logger.async_log_failure_event(
kwargs=timeout_kwargs,
response_obj=None,
start_time=time.time(),
end_time=time.time() + 30,
)
latency_key = f"{model_group}_map"
cached_data = await test_cache.async_get_cache(key=latency_key)
latency_list = cached_data[deployment_id]["latency"]
assert len(latency_list) == max_size
# Timeout penalty is the newest entry.
assert (
latency_list[-1] == 1000.0
), f"Timeout penalty should be at end of list, got {latency_list[-1]}"
# Oldest normal entry (1.0) has been discarded.
tolerance = 0.1
for latency in latency_list[:-1]:
assert (
abs(latency - 1.0) > tolerance
), f"Oldest latency 1.0 should have been discarded, found {latency}"
def test_list_order_preserved_after_multiple_trims():
"""
After many trims, the list still holds the most recent `max_size` entries
in insertion order (oldest at index 0, newest at index -1).
"""
max_size = 3
test_cache = DualCache()
lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
)
model_group = "gpt-3.5-turbo"
deployment_id = "test-deployment"
kwargs = {
"litellm_params": {
"metadata": {
"model_group": model_group,
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
}
}
# Add 10 entries (7 more than max) to trigger multiple trims.
all_latencies = []
for i in range(10):
start_time = time.time()
response_obj = {"usage": {"total_tokens": 1, "completion_tokens": 1}}
expected_latency = float(i + 1)
end_time = start_time + expected_latency
all_latencies.append(expected_latency)
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_map"
cached_data = test_cache.get_cache(key=latency_key)
latency_list = cached_data[deployment_id]["latency"]
assert len(latency_list) == max_size
# After inserting 1..10 with max_size=3, the list should be [8, 9, 10].
expected_remaining = all_latencies[-max_size:]
tolerance = 0.1
for i, expected in enumerate(expected_remaining):
assert (
abs(latency_list[i] - expected) < tolerance
), f"At index {i}, expected ~{expected}, got {latency_list[i]}"
@pytest.mark.asyncio
async def test_ttft_list_trimming_discards_oldest_entry_async():
"""
Async counterpart: the time_to_first_token list trims the oldest entry
when full. Exercises the async_log_success_event TTFT path, which only
runs when response_obj is a ModelResponse and the call is marked as
streaming with a completion_start_time.
"""
max_size = 3
test_cache = DualCache()
lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
)
model_group = "gpt-3.5-turbo"
deployment_id = "test-deployment"
ttft_values = []
for i in range(max_size + 1):
start_time = time.time()
expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4
completion_start_time = start_time + expected_ttft
end_time = start_time + float(i + 1)
ttft_values.append(expected_ttft)
kwargs = {
"litellm_params": {
"metadata": {
"model_group": model_group,
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
},
"stream": True,
"completion_start_time": completion_start_time,
}
response_obj = litellm.ModelResponse(
usage=litellm.Usage(completion_tokens=1, total_tokens=1)
)
await lowest_latency_logger.async_log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_map"
cached_data = await test_cache.async_get_cache(key=latency_key)
ttft_list = cached_data[deployment_id].get("time_to_first_token", [])
assert (
len(ttft_list) == max_size
), f"Expected {max_size} entries, got {len(ttft_list)}"
newest_ttft = ttft_values[-1]
oldest_ttft = ttft_values[0]
tolerance = 0.05
assert (
abs(ttft_list[-1] - newest_ttft) < tolerance
), f"Newest TTFT {newest_ttft} should be at end of list"
for ttft in ttft_list:
assert (
abs(ttft - oldest_ttft) > tolerance
), f"Oldest TTFT {oldest_ttft} should have been discarded"
Reference in a new issue View git blame Copy permalink