scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88
litellm/tests/local_testing/test_lowest_latency_routing.py
Clement 95c0f9db7d
fix(router): rank streaming latency routing by raw TTFT, not TTFT per token (#40202) 
* fix(router): rank streaming latency routing by raw TTFT, not TTFT per token

Latency-based routing divided time-to-first-token by completion_tokens
before storing it, so a deployment that streamed a long answer looked
faster to first token than one that answered briefly. TTFT is now stored
as plain seconds (first token time minus request start) in both the sync
and async success handlers, which is what the routing decision compares.

Non-streaming latency normalization per output token is unchanged.

Claude-Session: https://claude.ai/code/session_01Ttd5Q9ZhRPB4ch5guos3rj

* fix(router): store streaming TTFT under a seconds-only cache key

Workers on the previous release keep writing seconds-per-token samples
under "time_to_first_token" in the shared router cache during a rolling
deploy, so mixing the new raw-seconds samples into the same list averaged
incompatible units. Raw TTFT now lives under "time_to_first_token_seconds"
and routing reads only that key.

Also fix the regression test's token counts: with 50 tokens on the fast
deployment and 500 on the slow one the old per-token formula picks the
slow deployment, so the routing assertion now catches the bug.

Claude-Session: https://claude.ai/code/session_01Ttd5Q9ZhRPB4ch5guos3rj

* test(router): cover the TTFT sliding window from the unit-test shard

Move the TTFT list trimming checks from the CircleCI-only suite into the
mapped unit test file as one sync/async parametrized test, so the changed
lines in lowest_latency.py are exercised by the GitHub unit-test shard
that reports patch coverage.

Claude-Session: https://claude.ai/code/session_01Ttd5Q9ZhRPB4ch5guos3rj
2026-09-08 10:05:23 -07:00

1204 lines
37 KiB
Python
Raw 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"
@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]}"
Reference in a new issue View git blame Copy permalink