mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
update tests, uv black
This commit is contained in:
Oles Lomako
parent
b7215be349
commit
6beb2287c7
3 changed files with 636 additions and 505 deletions
|
|
@ -829,19 +829,21 @@ class Router:
|
|||
self._ensure_routing_strategy_logger(routing_strategy, routing_strategy_args)
|
||||
|
||||
def _ensure_routing_strategy_logger(
|
||||
self, routing_strategy: Union[RoutingStrategy, str], routing_strategy_args: Optional[dict] = None
|
||||
self,
|
||||
routing_strategy: Union[RoutingStrategy, str],
|
||||
routing_strategy_args: Optional[dict] = None,
|
||||
):
|
||||
"""
|
||||
Lazy initialization of routing strategy loggers.
|
||||
Creates logger if it doesn't exist yet. Supports per-request routing strategy overrides.
|
||||
"""
|
||||
routing_strategy_args = routing_strategy_args or {}
|
||||
|
||||
|
||||
if (
|
||||
routing_strategy == RoutingStrategy.LEAST_BUSY.value
|
||||
or routing_strategy == RoutingStrategy.LEAST_BUSY
|
||||
):
|
||||
if not hasattr(self, 'leastbusy_logger') or self.leastbusy_logger is None:
|
||||
if not hasattr(self, "leastbusy_logger") or self.leastbusy_logger is None:
|
||||
self.leastbusy_logger = LeastBusyLoggingHandler(router_cache=self.cache)
|
||||
## add callback
|
||||
if isinstance(litellm.input_callback, list):
|
||||
|
|
@ -854,7 +856,7 @@ class Router:
|
|||
routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING.value
|
||||
or routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING
|
||||
):
|
||||
if not hasattr(self, 'lowesttpm_logger') or self.lowesttpm_logger is None:
|
||||
if not hasattr(self, "lowesttpm_logger") or self.lowesttpm_logger is None:
|
||||
self.lowesttpm_logger = LowestTPMLoggingHandler(
|
||||
router_cache=self.cache,
|
||||
routing_args=routing_strategy_args,
|
||||
|
|
@ -865,7 +867,10 @@ class Router:
|
|||
routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING_V2.value
|
||||
or routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING_V2
|
||||
):
|
||||
if not hasattr(self, 'lowesttpm_logger_v2') or self.lowesttpm_logger_v2 is None:
|
||||
if (
|
||||
not hasattr(self, "lowesttpm_logger_v2")
|
||||
or self.lowesttpm_logger_v2 is None
|
||||
):
|
||||
self.lowesttpm_logger_v2 = LowestTPMLoggingHandler_v2(
|
||||
router_cache=self.cache,
|
||||
routing_args=routing_strategy_args,
|
||||
|
|
@ -876,7 +881,10 @@ class Router:
|
|||
routing_strategy == RoutingStrategy.LATENCY_BASED.value
|
||||
or routing_strategy == RoutingStrategy.LATENCY_BASED
|
||||
):
|
||||
if not hasattr(self, 'lowestlatency_logger') or self.lowestlatency_logger is None:
|
||||
if (
|
||||
not hasattr(self, "lowestlatency_logger")
|
||||
or self.lowestlatency_logger is None
|
||||
):
|
||||
self.lowestlatency_logger = LowestLatencyLoggingHandler(
|
||||
router_cache=self.cache,
|
||||
routing_args=routing_strategy_args,
|
||||
|
|
@ -887,7 +895,7 @@ class Router:
|
|||
routing_strategy == RoutingStrategy.COST_BASED.value
|
||||
or routing_strategy == RoutingStrategy.COST_BASED
|
||||
):
|
||||
if not hasattr(self, 'lowestcost_logger') or self.lowestcost_logger is None:
|
||||
if not hasattr(self, "lowestcost_logger") or self.lowestcost_logger is None:
|
||||
self.lowestcost_logger = LowestCostLoggingHandler(
|
||||
router_cache=self.cache,
|
||||
routing_args={},
|
||||
|
|
@ -9506,8 +9514,10 @@ class Router:
|
|||
|
||||
start_time = time.time()
|
||||
# Ensure logger is initialized for the routing strategy (lazy init)
|
||||
self._ensure_routing_strategy_logger(routing_strategy_to_use, self.routing_strategy_args)
|
||||
|
||||
self._ensure_routing_strategy_logger(
|
||||
routing_strategy_to_use, self.routing_strategy_args
|
||||
)
|
||||
|
||||
if (
|
||||
routing_strategy_to_use == "usage-based-routing-v2"
|
||||
and self.lowesttpm_logger_v2 is not None
|
||||
|
|
@ -9677,7 +9687,9 @@ class Router:
|
|||
) # Pass-through uses global routing strategy
|
||||
|
||||
# Ensure logger is initialized (lazy init)
|
||||
self._ensure_routing_strategy_logger(routing_strategy_to_use, self.routing_strategy_args)
|
||||
self._ensure_routing_strategy_logger(
|
||||
routing_strategy_to_use, self.routing_strategy_args
|
||||
)
|
||||
|
||||
if (
|
||||
routing_strategy_to_use == "usage-based-routing-v2"
|
||||
|
|
@ -9830,7 +9842,9 @@ class Router:
|
|||
|
||||
# Support per-request routing_strategy override from key/team config
|
||||
if request_kwargs is not None:
|
||||
routing_strategy_to_use = request_kwargs.pop("routing_strategy", None) or self.routing_strategy
|
||||
routing_strategy_to_use = (
|
||||
request_kwargs.pop("routing_strategy", None) or self.routing_strategy
|
||||
)
|
||||
else:
|
||||
routing_strategy_to_use = self.routing_strategy
|
||||
|
||||
|
|
@ -9903,8 +9917,10 @@ class Router:
|
|||
)
|
||||
|
||||
# Ensure logger is initialized for the routing strategy (lazy init)
|
||||
self._ensure_routing_strategy_logger(routing_strategy_to_use, self.routing_strategy_args)
|
||||
|
||||
self._ensure_routing_strategy_logger(
|
||||
routing_strategy_to_use, self.routing_strategy_args
|
||||
)
|
||||
|
||||
if (
|
||||
routing_strategy_to_use == "least-busy"
|
||||
and self.leastbusy_logger is not None
|
||||
|
|
|
|||
|
|
@ -2306,3 +2306,274 @@ def test_get_router_model_info_with_deployment_object():
|
|||
# Verify we got valid model info back
|
||||
assert model_info is not None
|
||||
assert isinstance(model_info, dict)
|
||||
|
||||
|
||||
class TestEnsureRoutingStrategyLogger:
|
||||
"""Test the _ensure_routing_strategy_logger method for lazy initialization."""
|
||||
|
||||
def test_lazy_initialization_latency_logger(self):
|
||||
"""Test that latency logger is created on-demand when it doesn't exist."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "fake-key-1",
|
||||
"rpm": 100,
|
||||
"tpm": 10000,
|
||||
},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="simple-shuffle", # Don't initialize latency logger
|
||||
)
|
||||
|
||||
# Verify logger doesn't exist initially
|
||||
assert (
|
||||
not hasattr(router, "lowestlatency_logger")
|
||||
or router.lowestlatency_logger is None
|
||||
)
|
||||
|
||||
# Call _ensure_routing_strategy_logger for latency-based-routing
|
||||
router._ensure_routing_strategy_logger("latency-based-routing", {"ttl": 1})
|
||||
|
||||
# Verify logger was created
|
||||
assert (
|
||||
hasattr(router, "lowestlatency_logger")
|
||||
and router.lowestlatency_logger is not None
|
||||
)
|
||||
|
||||
def test_lazy_initialization_cost_logger(self):
|
||||
"""Test that cost logger is created on-demand when it doesn't exist."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "fake-key-1",
|
||||
"rpm": 100,
|
||||
"tpm": 10000,
|
||||
},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="simple-shuffle", # Don't initialize cost logger
|
||||
)
|
||||
|
||||
# Verify logger doesn't exist initially
|
||||
assert (
|
||||
not hasattr(router, "lowestcost_logger") or router.lowestcost_logger is None
|
||||
)
|
||||
|
||||
# Call _ensure_routing_strategy_logger for cost-based-routing
|
||||
router._ensure_routing_strategy_logger("cost-based-routing")
|
||||
|
||||
# Verify logger was created
|
||||
assert (
|
||||
hasattr(router, "lowestcost_logger")
|
||||
and router.lowestcost_logger is not None
|
||||
)
|
||||
|
||||
def test_lazy_initialization_tpm_logger_v2(self):
|
||||
"""Test that TPM logger v2 is created on-demand when it doesn't exist."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "fake-key-1",
|
||||
"rpm": 100,
|
||||
"tpm": 10000,
|
||||
},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="simple-shuffle", # Don't initialize TPM logger
|
||||
)
|
||||
|
||||
# Verify logger doesn't exist initially
|
||||
assert (
|
||||
not hasattr(router, "lowesttpm_logger_v2")
|
||||
or router.lowesttpm_logger_v2 is None
|
||||
)
|
||||
|
||||
# Call _ensure_routing_strategy_logger for usage-based-routing-v2
|
||||
router._ensure_routing_strategy_logger("usage-based-routing-v2")
|
||||
|
||||
# Verify logger was created
|
||||
assert (
|
||||
hasattr(router, "lowesttpm_logger_v2")
|
||||
and router.lowesttpm_logger_v2 is not None
|
||||
)
|
||||
|
||||
def test_lazy_initialization_least_busy_logger(self):
|
||||
"""Test that least-busy logger is created on-demand when it doesn't exist."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "fake-key-1",
|
||||
"rpm": 100,
|
||||
"tpm": 10000,
|
||||
},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="simple-shuffle", # Don't initialize least-busy logger
|
||||
)
|
||||
|
||||
# Verify logger doesn't exist initially
|
||||
assert (
|
||||
not hasattr(router, "leastbusy_logger") or router.leastbusy_logger is None
|
||||
)
|
||||
|
||||
# Call _ensure_routing_strategy_logger for least-busy
|
||||
router._ensure_routing_strategy_logger("least-busy")
|
||||
|
||||
# Verify logger was created
|
||||
assert (
|
||||
hasattr(router, "leastbusy_logger") and router.leastbusy_logger is not None
|
||||
)
|
||||
|
||||
def test_no_reinitialize_existing_logger(self):
|
||||
"""Test that existing loggers are not re-initialized."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "fake-key-1",
|
||||
"rpm": 100,
|
||||
"tpm": 10000,
|
||||
},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="latency-based-routing", # Initialize latency logger
|
||||
routing_strategy_args={"ttl": 1},
|
||||
)
|
||||
|
||||
# Get reference to the existing logger
|
||||
original_logger = router.lowestlatency_logger
|
||||
assert original_logger is not None
|
||||
|
||||
# Call _ensure_routing_strategy_logger again
|
||||
router._ensure_routing_strategy_logger("latency-based-routing", {"ttl": 1})
|
||||
|
||||
# Verify the same logger instance is still being used (not re-initialized)
|
||||
assert router.lowestlatency_logger is original_logger
|
||||
|
||||
def test_no_reinitialize_cost_logger(self):
|
||||
"""Test that existing cost logger is not re-initialized."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "fake-key-1",
|
||||
"rpm": 100,
|
||||
"tpm": 10000,
|
||||
},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="cost-based-routing", # Initialize cost logger
|
||||
)
|
||||
|
||||
# Get reference to the existing logger
|
||||
original_logger = router.lowestcost_logger
|
||||
assert original_logger is not None
|
||||
|
||||
# Call _ensure_routing_strategy_logger again
|
||||
router._ensure_routing_strategy_logger("cost-based-routing")
|
||||
|
||||
# Verify the same logger instance is still being used (not re-initialized)
|
||||
assert router.lowestcost_logger is original_logger
|
||||
|
||||
def test_simple_shuffle_no_logger(self):
|
||||
"""Test that simple-shuffle doesn't create any logger."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "fake-key-1",
|
||||
"rpm": 100,
|
||||
"tpm": 10000,
|
||||
},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Call _ensure_routing_strategy_logger for simple-shuffle
|
||||
router._ensure_routing_strategy_logger("simple-shuffle")
|
||||
|
||||
# Verify no routing strategy loggers were created
|
||||
assert (
|
||||
not hasattr(router, "lowestlatency_logger")
|
||||
or router.lowestlatency_logger is None
|
||||
)
|
||||
assert (
|
||||
not hasattr(router, "lowestcost_logger") or router.lowestcost_logger is None
|
||||
)
|
||||
assert (
|
||||
not hasattr(router, "lowesttpm_logger_v2")
|
||||
or router.lowesttpm_logger_v2 is None
|
||||
)
|
||||
assert (
|
||||
not hasattr(router, "leastbusy_logger") or router.leastbusy_logger is None
|
||||
)
|
||||
|
||||
def test_multiple_loggers_initialization(self):
|
||||
"""Test that multiple different loggers can be initialized."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "fake-key-1",
|
||||
"rpm": 100,
|
||||
"tpm": 10000,
|
||||
},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
]
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Initialize multiple loggers
|
||||
router._ensure_routing_strategy_logger("latency-based-routing", {"ttl": 1})
|
||||
router._ensure_routing_strategy_logger("cost-based-routing")
|
||||
router._ensure_routing_strategy_logger("least-busy")
|
||||
|
||||
# Verify all loggers were created
|
||||
assert (
|
||||
hasattr(router, "lowestlatency_logger")
|
||||
and router.lowestlatency_logger is not None
|
||||
)
|
||||
assert (
|
||||
hasattr(router, "lowestcost_logger")
|
||||
and router.lowestcost_logger is not None
|
||||
)
|
||||
assert (
|
||||
hasattr(router, "leastbusy_logger") and router.leastbusy_logger is not None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,12 +1,12 @@
|
|||
"""
|
||||
Unit tests for routing strategy override functionality.
|
||||
|
||||
Tests cover:
|
||||
1. Per-request routing strategy overrides with lazy logger initialization
|
||||
2. Regression fixes for routing strategy handling
|
||||
3. Baseline functionality for each routing strategy
|
||||
4. Fallback behavior when routing strategies are overridden
|
||||
Tests verify:
|
||||
1. Algorithm correctness - each strategy selects the expected deployment
|
||||
2. Per-request routing strategy overrides work correctly
|
||||
3. Critical regression fixes (NameError prevention, override preservation)
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import pytest
|
||||
|
|
@ -15,6 +15,7 @@ from unittest.mock import patch
|
|||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
from litellm import Router
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -85,178 +86,12 @@ def pass_through_model_list():
|
|||
]
|
||||
|
||||
|
||||
class TestRoutingStrategyOverride:
|
||||
"""Test per-request routing strategy overrides with lazy initialization."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_override_simple_shuffle_to_latency_based(self, base_model_list):
|
||||
"""Test per-request override from simple-shuffle to latency-based-routing."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
routing_strategy_args={"ttl": 1},
|
||||
)
|
||||
|
||||
# Test that override triggers lazy initialization and works
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={"routing_strategy": "latency-based-routing"},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Verify deployment is selected and latency logger was lazily initialized
|
||||
assert deployment is not None
|
||||
assert hasattr(router, 'lowestlatency_logger') and router.lowestlatency_logger is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_override_to_cost_based_routing(self, base_model_list):
|
||||
"""Test per-request override to cost-based-routing."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Test that override triggers lazy initialization and works
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={"routing_strategy": "cost-based-routing"},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Verify deployment is selected and cost logger was lazily initialized
|
||||
assert deployment is not None
|
||||
assert hasattr(router, 'lowestcost_logger') and router.lowestcost_logger is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_override_to_usage_based_routing_v2(self, base_model_list):
|
||||
"""Test per-request override to usage-based-routing-v2."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Test that override triggers lazy initialization and works
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={"routing_strategy": "usage-based-routing-v2"},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Verify deployment is selected and TPM logger was lazily initialized
|
||||
assert deployment is not None
|
||||
assert hasattr(router, 'lowesttpm_logger_v2') and router.lowesttpm_logger_v2 is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_override_to_least_busy(self, base_model_list):
|
||||
"""Test per-request override to least-busy."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Test that override triggers lazy initialization and works
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={"routing_strategy": "least-busy"},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Verify deployment is selected and least-busy logger was lazily initialized
|
||||
assert deployment is not None
|
||||
assert hasattr(router, 'leastbusy_logger') and router.leastbusy_logger is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_with_cost_based_routing_no_crash(
|
||||
self, pass_through_model_list
|
||||
):
|
||||
"""
|
||||
Test pass-through with cost-based-routing doesn't crash.
|
||||
Tests that deployment = None initialization prevents NameError.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=pass_through_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# This should not crash even with cost-based-routing (lazy init will create logger)
|
||||
# Note: pass-through uses global routing_strategy, not per-request override (for now)
|
||||
deployment = (
|
||||
await router.async_get_available_deployment_for_pass_through(
|
||||
model="gpt-4",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
)
|
||||
|
||||
assert deployment is not None
|
||||
assert deployment.get("litellm_params", {}).get("use_in_pass_through") is True
|
||||
|
||||
def test_latency_override_in_sync_path(self, base_model_list):
|
||||
"""
|
||||
Test that latency override works in sync path.
|
||||
Tests routing_strategy_to_use variable is used instead of self.routing_strategy.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
routing_strategy_args={"ttl": 1},
|
||||
)
|
||||
|
||||
# Test latency-based routing with override (lazy init will create logger)
|
||||
deployment = router.get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={"routing_strategy": "latency-based-routing"},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Verify deployment is selected and logger was lazily initialized
|
||||
assert deployment is not None
|
||||
assert hasattr(router, 'lowestlatency_logger') and router.lowestlatency_logger is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_to_sync_fallthrough_preserves_override(self, base_model_list):
|
||||
"""
|
||||
Test async→sync fallthrough preserves routing_strategy override.
|
||||
Tests that override is preserved in request_kwargs.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Mock get_available_deployment to capture the routing_strategy passed
|
||||
captured_kwargs = {}
|
||||
|
||||
def capture_kwargs(*args, **kwargs):
|
||||
captured_kwargs.update(kwargs.get("request_kwargs", {}))
|
||||
# Return a valid deployment
|
||||
return base_model_list[0]
|
||||
|
||||
with patch.object(
|
||||
router, "get_available_deployment", side_effect=capture_kwargs
|
||||
):
|
||||
# Trigger the async→sync fallthrough path by using an unsupported routing strategy
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={
|
||||
"routing_strategy": "invalid-unknown-strategy" # Not in the async whitelist
|
||||
},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Verify routing_strategy was preserved in request_kwargs
|
||||
assert (
|
||||
"routing_strategy" in captured_kwargs
|
||||
), "routing_strategy should be preserved in fallthrough"
|
||||
assert captured_kwargs["routing_strategy"] == "invalid-unknown-strategy"
|
||||
|
||||
|
||||
class TestCostBasedRoutingBaseline:
|
||||
"""Test baseline functionality for cost-based routing."""
|
||||
class TestRoutingAlgorithmCorrectness:
|
||||
"""Test that each routing strategy selects the correct deployment."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cost_based_routing_selects_cheapest(self):
|
||||
"""Test that cost-based routing selects the cheapest deployment."""
|
||||
"""Verify cost-based routing picks the cheapest deployment."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "test-model",
|
||||
|
|
@ -273,11 +108,21 @@ class TestCostBasedRoutingBaseline:
|
|||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "key2",
|
||||
"input_cost_per_token": 0.001,
|
||||
"input_cost_per_token": 0.001, # ← Should win (cheapest)
|
||||
"output_cost_per_token": 0.001,
|
||||
},
|
||||
"model_info": {"id": "cheap-1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "key3",
|
||||
"input_cost_per_token": 0.0015,
|
||||
"output_cost_per_token": 0.0015,
|
||||
},
|
||||
"model_info": {"id": "medium-1"},
|
||||
},
|
||||
]
|
||||
|
||||
router = Router(
|
||||
|
|
@ -294,174 +139,303 @@ class TestCostBasedRoutingBaseline:
|
|||
# Should select the cheaper deployment
|
||||
assert deployment["model_info"]["id"] == "cheap-1"
|
||||
|
||||
@pytest.mark.skip(reason="Cost-based routing doesn't have sync implementation yet")
|
||||
def test_cost_based_routing_sync_selects_cheapest(self):
|
||||
"""Test that cost-based routing works in sync path."""
|
||||
pass
|
||||
|
||||
|
||||
class TestGlobalRoutingStrategy:
|
||||
"""Test that global routing strategy still works when no override provided."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_latency_based_routing_async(self, base_model_list):
|
||||
"""Test global latency-based routing works without override."""
|
||||
async def test_latency_based_selects_fastest_deployment(self):
|
||||
"""Verify latency-based routing picks deployment with lowest latency."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key1"},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key2"},
|
||||
"model_info": {"id": "deployment-2"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key3"},
|
||||
"model_info": {"id": "deployment-3"},
|
||||
},
|
||||
]
|
||||
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
model_list=model_list,
|
||||
routing_strategy="latency-based-routing",
|
||||
routing_strategy_args={"ttl": 1},
|
||||
)
|
||||
|
||||
# Pre-populate latency cache:
|
||||
# deployment-1: 100ms latency
|
||||
# deployment-2: 50ms latency ← Should win (fastest)
|
||||
# deployment-3: 200ms latency
|
||||
latency_cache = {
|
||||
"deployment-1": {"latency": [0.1, 0.1, 0.1]}, # 100ms avg
|
||||
"deployment-2": {"latency": [0.05, 0.05, 0.05]}, # 50ms avg - fastest
|
||||
"deployment-3": {"latency": [0.2, 0.2, 0.2]}, # 200ms avg
|
||||
}
|
||||
router.cache.set_cache(key="test-model_map", value=latency_cache)
|
||||
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={}, # No override
|
||||
model="test-model",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
assert deployment is not None
|
||||
assert router.lowestlatency_logger is not None
|
||||
|
||||
def test_global_latency_based_routing_sync(self, base_model_list):
|
||||
"""Test global latency-based routing works in sync path without override."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="latency-based-routing",
|
||||
routing_strategy_args={"ttl": 1},
|
||||
)
|
||||
|
||||
deployment = router.get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={}, # No override
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
assert deployment is not None
|
||||
assert router.lowestlatency_logger is not None
|
||||
# Should select deployment with lowest latency
|
||||
assert deployment["model_info"]["id"] == "deployment-2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_usage_based_routing_v2(self, base_model_list):
|
||||
"""Test global usage-based-routing-v2 works without override."""
|
||||
async def test_usage_based_v2_selects_most_available_tpm(self):
|
||||
"""Verify usage-based-v2 picks deployment with most available TPM."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "key1",
|
||||
"tpm": 100,
|
||||
},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "key2",
|
||||
"tpm": 100,
|
||||
},
|
||||
"model_info": {"id": "deployment-2"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "key3",
|
||||
"tpm": 100,
|
||||
},
|
||||
"model_info": {"id": "deployment-3"},
|
||||
},
|
||||
]
|
||||
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
model_list=model_list,
|
||||
routing_strategy="usage-based-routing-v2",
|
||||
)
|
||||
|
||||
# Pre-populate usage cache (current minute format):
|
||||
# deployment-1: 50/100 TPM used
|
||||
# deployment-2: 20/100 TPM used ← Should win (80 available)
|
||||
# deployment-3: 90/100 TPM used
|
||||
dt = get_utc_datetime()
|
||||
current_minute = dt.strftime("%H-%M")
|
||||
router.cache.set_cache(
|
||||
key=f"deployment-1:gpt-3.5-turbo:tpm:{current_minute}", value=50
|
||||
)
|
||||
router.cache.set_cache(
|
||||
key=f"deployment-2:gpt-3.5-turbo:tpm:{current_minute}", value=20
|
||||
) # Most available
|
||||
router.cache.set_cache(
|
||||
key=f"deployment-3:gpt-3.5-turbo:tpm:{current_minute}", value=90
|
||||
)
|
||||
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={}, # No override
|
||||
model="test-model",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
assert deployment is not None
|
||||
assert router.lowesttpm_logger_v2 is not None
|
||||
# Should select deployment with most available TPM
|
||||
assert deployment["model_info"]["id"] == "deployment-2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_least_busy(self, base_model_list):
|
||||
"""Test global least-busy works without override."""
|
||||
async def test_least_busy_selects_deployment_with_fewest_requests(self):
|
||||
"""Verify least-busy picks deployment with fewest active requests."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key1"},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key2"},
|
||||
"model_info": {"id": "deployment-2"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key3"},
|
||||
"model_info": {"id": "deployment-3"},
|
||||
},
|
||||
]
|
||||
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
model_list=model_list,
|
||||
routing_strategy="least-busy",
|
||||
)
|
||||
|
||||
# Pre-populate request count cache:
|
||||
# deployment-1: 5 active requests
|
||||
# deployment-2: 2 active requests ← Should win (least busy)
|
||||
# deployment-3: 8 active requests
|
||||
request_count_cache = {
|
||||
"deployment-1": 5,
|
||||
"deployment-2": 2, # Least busy
|
||||
"deployment-3": 8,
|
||||
}
|
||||
router.cache.set_cache(
|
||||
key="test-model_request_count", value=request_count_cache
|
||||
)
|
||||
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={}, # No override
|
||||
model="test-model",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
assert deployment is not None
|
||||
assert router.leastbusy_logger is not None
|
||||
# Should select deployment with fewest active requests
|
||||
assert deployment["model_info"]["id"] == "deployment-2"
|
||||
|
||||
|
||||
class TestOverridePrecedence:
|
||||
"""Test that request kwargs override takes precedence over global settings."""
|
||||
class TestRequestBodyRoutingStrategyOverride:
|
||||
"""Test routing_strategy override passed in request body actually uses the correct algorithm."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_override_takes_precedence_async(self, base_model_list):
|
||||
"""Test request override takes precedence over global routing strategy."""
|
||||
async def test_request_body_override_actually_uses_cost_based(self):
|
||||
"""Verify request body cost-based override ACTUALLY selects cheapest."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "key1",
|
||||
"input_cost_per_token": 0.002, # Expensive
|
||||
"output_cost_per_token": 0.002,
|
||||
},
|
||||
"model_info": {"id": "expensive-1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "key2",
|
||||
"input_cost_per_token": 0.001, # ← Should win (cheapest)
|
||||
"output_cost_per_token": 0.001,
|
||||
},
|
||||
"model_info": {"id": "cheap-1"},
|
||||
},
|
||||
]
|
||||
|
||||
# Global: simple-shuffle
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="least-busy", # Global strategy
|
||||
model_list=model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Override to latency-based (lazy init will create logger)
|
||||
# Request override: cost-based-routing
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={
|
||||
"routing_strategy": "latency-based-routing" # Override
|
||||
},
|
||||
model="test-model",
|
||||
request_kwargs={"routing_strategy": "cost-based-routing"},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Verify override worked and latency logger was lazily initialized
|
||||
assert deployment is not None
|
||||
assert hasattr(router, 'lowestlatency_logger') and router.lowestlatency_logger is not None
|
||||
|
||||
def test_override_takes_precedence_sync(self, base_model_list):
|
||||
"""Test request override takes precedence in sync path."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="least-busy", # Global strategy
|
||||
)
|
||||
|
||||
# Override to latency-based (lazy init will create logger)
|
||||
deployment = router.get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={
|
||||
"routing_strategy": "latency-based-routing" # Override
|
||||
},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Verify override worked and latency logger was lazily initialized
|
||||
assert deployment is not None
|
||||
assert hasattr(router, 'lowestlatency_logger') and router.lowestlatency_logger is not None
|
||||
# Verify cheap deployment is selected (not random from simple-shuffle)
|
||||
assert deployment["model_info"]["id"] == "cheap-1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_overrides_different_requests(self, base_model_list):
|
||||
"""Test that different requests can use different routing strategies."""
|
||||
async def test_request_body_override_actually_uses_latency_based(self):
|
||||
"""Verify request body latency-based override ACTUALLY selects fastest."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key1"},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key2"},
|
||||
"model_info": {"id": "deployment-2"},
|
||||
},
|
||||
]
|
||||
|
||||
# Global: simple-shuffle
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Pre-populate latency cache
|
||||
latency_cache = {
|
||||
"deployment-1": {"latency": [0.2, 0.2]}, # 200ms avg
|
||||
"deployment-2": {"latency": [0.05, 0.05]}, # 50ms avg - fastest
|
||||
}
|
||||
router.cache.set_cache(key="test-model_map", value=latency_cache)
|
||||
|
||||
# Request override: latency-based-routing
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="test-model",
|
||||
request_kwargs={"routing_strategy": "latency-based-routing"},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Verify fastest deployment is selected (not random from simple-shuffle)
|
||||
assert deployment["model_info"]["id"] == "deployment-2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_body_routing_strategy_is_popped(self, base_model_list):
|
||||
"""
|
||||
Test that routing_strategy is popped from request_kwargs after use.
|
||||
|
||||
Ensures routing_strategy doesn't leak to downstream LLM provider calls.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Request 1: Use cost-based (lazy init will create logger)
|
||||
deployment1 = await router.async_get_available_deployment(
|
||||
request_kwargs = {
|
||||
"routing_strategy": "latency-based-routing",
|
||||
"temperature": 0.7, # Other param that should remain
|
||||
}
|
||||
|
||||
# Simulate curl request body
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={"routing_strategy": "cost-based-routing"},
|
||||
messages=[{"role": "user", "content": "test1"}],
|
||||
request_kwargs=request_kwargs,
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Request 2: Use latency-based (lazy init will create logger)
|
||||
deployment2 = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={"routing_strategy": "latency-based-routing"},
|
||||
messages=[{"role": "user", "content": "test2"}],
|
||||
# Verify deployment is selected
|
||||
assert deployment is not None
|
||||
|
||||
# Verify routing_strategy was popped from request_kwargs
|
||||
assert "routing_strategy" not in request_kwargs, (
|
||||
"routing_strategy should be popped from request_kwargs "
|
||||
"to prevent it from being passed to LLM provider"
|
||||
)
|
||||
|
||||
# Both deployments should be valid
|
||||
assert deployment1 is not None
|
||||
assert deployment2 is not None
|
||||
# Verify both loggers were lazily initialized
|
||||
assert hasattr(router, 'lowestcost_logger') and router.lowestcost_logger is not None
|
||||
assert hasattr(router, 'lowestlatency_logger') and router.lowestlatency_logger is not None
|
||||
# Verify other params remain
|
||||
assert "temperature" in request_kwargs
|
||||
assert request_kwargs["temperature"] == 0.7
|
||||
|
||||
|
||||
class TestRegressionFixes:
|
||||
"""Test specific regression fixes."""
|
||||
"""Test specific regression fixes for critical bugs."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_cost_routing_no_name_error(self, pass_through_model_list):
|
||||
async def test_pass_through_cost_routing_no_name_error(
|
||||
self, pass_through_model_list
|
||||
):
|
||||
"""
|
||||
Test that deployment = None initialization prevents NameError.
|
||||
Uses simple-shuffle to avoid cost-based routing issues in pass-through.
|
||||
Regression test for bug where uninitialized deployment variable caused NameError.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=pass_through_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# This should not raise NameError (our fix adds deployment = None initialization)
|
||||
# This should not raise NameError (fix adds deployment = None initialization)
|
||||
deployment = await router.async_get_available_deployment_for_pass_through(
|
||||
model="gpt-4",
|
||||
request_kwargs={},
|
||||
|
|
@ -474,7 +448,7 @@ class TestRegressionFixes:
|
|||
def test_sync_latency_respects_per_request_override(self, base_model_list):
|
||||
"""
|
||||
Test that sync latency-based routing uses per-request override, not global strategy.
|
||||
Verifies routing_strategy_to_use variable is used instead of self.routing_strategy.
|
||||
Regression test: verifies routing_strategy_to_use variable is used instead of self.routing_strategy.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
|
|
@ -491,13 +465,16 @@ class TestRegressionFixes:
|
|||
|
||||
# Verify override worked and latency logger was lazily initialized
|
||||
assert deployment is not None
|
||||
assert hasattr(router, 'lowestlatency_logger') and router.lowestlatency_logger is not None
|
||||
assert (
|
||||
hasattr(router, "lowestlatency_logger")
|
||||
and router.lowestlatency_logger is not None
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_to_sync_fallthrough_preserves_override(self, base_model_list):
|
||||
"""
|
||||
Test that routing_strategy override is preserved when async falls through to sync.
|
||||
Verifies the override is not lost in the fallthrough path.
|
||||
Regression test: verifies the override is not lost in the fallthrough path.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
|
|
@ -519,7 +496,7 @@ class TestRegressionFixes:
|
|||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={
|
||||
"routing_strategy": "invalid-unknown-strategy" # Not in the async whitelist
|
||||
"routing_strategy": "invalid-unknown-strategy" # Not in async whitelist
|
||||
},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
|
@ -530,203 +507,12 @@ class TestRegressionFixes:
|
|||
), "routing_strategy should be preserved in fallthrough"
|
||||
assert captured_kwargs["routing_strategy"] == "invalid-unknown-strategy"
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
"""Test edge cases and error conditions."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_routing_strategy_override(self, base_model_list):
|
||||
"""Test that invalid routing strategy override is handled gracefully."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Invalid routing strategy should fallback to sync path
|
||||
with patch.object(
|
||||
router, "get_available_deployment", return_value=base_model_list[0]
|
||||
) as mock_sync:
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={"routing_strategy": "invalid-strategy"},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Should fallback to sync get_available_deployment
|
||||
mock_sync.assert_called_once()
|
||||
assert deployment is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_request_kwargs_with_override(self, base_model_list):
|
||||
"""Test routing strategy override when request_kwargs is None."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Should use global routing strategy when request_kwargs is None
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
assert deployment is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_routing_strategy_override(self, base_model_list):
|
||||
"""Test that empty routing_strategy in request_kwargs uses global."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="latency-based-routing",
|
||||
routing_strategy_args={"ttl": 1},
|
||||
)
|
||||
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={"routing_strategy": None}, # Empty override
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Should use global strategy (latency-based)
|
||||
assert deployment is not None
|
||||
assert router.lowestlatency_logger is not None
|
||||
|
||||
|
||||
class TestPassThroughEndpoints:
|
||||
"""Test routing strategy overrides with pass-through endpoints."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_respects_override(self, pass_through_model_list):
|
||||
"""Test that pass-through endpoints work with different routing strategies."""
|
||||
router = Router(
|
||||
model_list=pass_through_model_list,
|
||||
routing_strategy="least-busy", # Global uses least-busy
|
||||
)
|
||||
|
||||
# Pass-through currently uses global routing strategy
|
||||
deployment = (
|
||||
await router.async_get_available_deployment_for_pass_through(
|
||||
model="gpt-4",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
)
|
||||
|
||||
assert deployment is not None
|
||||
assert deployment.get("litellm_params", {}).get("use_in_pass_through") is True
|
||||
assert router.leastbusy_logger is not None
|
||||
|
||||
def test_pass_through_sync_with_override(self, pass_through_model_list):
|
||||
"""Test sync pass-through with routing strategy override."""
|
||||
router = Router(
|
||||
model_list=pass_through_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
# Note: Sync pass-through uses global routing strategy (not overridable yet)
|
||||
# This test verifies it doesn't crash
|
||||
deployment = router.get_available_deployment_for_pass_through(
|
||||
model="gpt-4",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
assert deployment is not None
|
||||
assert deployment.get("litellm_params", {}).get("use_in_pass_through") is True
|
||||
|
||||
|
||||
class TestSimpleShuffleOverride:
|
||||
"""Test simple-shuffle as both global and override."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_override_to_simple_shuffle(self, base_model_list):
|
||||
"""Test overriding to simple-shuffle from another strategy."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="latency-based-routing",
|
||||
routing_strategy_args={"ttl": 1},
|
||||
)
|
||||
|
||||
# Override to simple-shuffle
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={"routing_strategy": "simple-shuffle"},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Verify override worked
|
||||
assert deployment is not None
|
||||
|
||||
def test_simple_shuffle_returns_valid_deployment(self, base_model_list):
|
||||
"""Test that simple-shuffle returns a valid deployment."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
deployment = router.get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Should return one of the valid deployments
|
||||
assert deployment is not None
|
||||
assert deployment["model_info"]["id"] in [
|
||||
"deployment-1",
|
||||
"deployment-2",
|
||||
"deployment-3",
|
||||
]
|
||||
|
||||
|
||||
class TestRoutingStrategyIntegration:
|
||||
"""Integration tests for routing strategies with overrides."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_routing_strategies_work(self, base_model_list):
|
||||
"""Test that all routing strategies can be used as overrides with lazy initialization."""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
strategies_to_test = [
|
||||
"simple-shuffle",
|
||||
"latency-based-routing",
|
||||
"cost-based-routing",
|
||||
"usage-based-routing-v2",
|
||||
"least-busy",
|
||||
]
|
||||
|
||||
for strategy in strategies_to_test:
|
||||
# Trigger lazy initialization by making request with override
|
||||
# This will create the logger if it doesn't exist
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs={"routing_strategy": strategy},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
assert deployment is not None, f"Failed for strategy: {strategy}"
|
||||
|
||||
# Verify the appropriate logger was created via lazy init
|
||||
if strategy == "latency-based-routing":
|
||||
assert hasattr(router, 'lowestlatency_logger') and router.lowestlatency_logger is not None
|
||||
elif strategy == "cost-based-routing":
|
||||
assert hasattr(router, 'lowestcost_logger') and router.lowestcost_logger is not None
|
||||
elif strategy == "usage-based-routing-v2":
|
||||
assert hasattr(router, 'lowesttpm_logger_v2') and router.lowesttpm_logger_v2 is not None
|
||||
elif strategy == "least-busy":
|
||||
assert hasattr(router, 'leastbusy_logger') and router.leastbusy_logger is not None
|
||||
|
||||
|
||||
class TestRoutingStrategyPersistence:
|
||||
"""Test that routing strategy doesn't leak between requests."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_strategy_not_mutated(self, base_model_list):
|
||||
"""Test that request_kwargs routing_strategy is properly popped and doesn't affect router state."""
|
||||
async def test_routing_strategy_not_mutated_across_requests(self, base_model_list):
|
||||
"""
|
||||
Test that routing_strategy override doesn't mutate router's global strategy.
|
||||
Regression test: ensures per-request override is truly per-request.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
|
|
@ -745,25 +531,83 @@ class TestRoutingStrategyPersistence:
|
|||
assert router.routing_strategy == original_strategy
|
||||
assert router.routing_strategy == "simple-shuffle"
|
||||
|
||||
|
||||
class TestOverridePrecedence:
|
||||
"""Test that per-request override takes precedence over global settings."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_kwargs_not_mutated_by_pop(self, base_model_list):
|
||||
"""Test that popping routing_strategy from request_kwargs doesn't affect original dict."""
|
||||
async def test_override_takes_precedence_over_global_async(self):
|
||||
"""Verify request override takes precedence over global routing strategy."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key1"},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key2"},
|
||||
"model_info": {"id": "deployment-2"},
|
||||
},
|
||||
]
|
||||
|
||||
# Global: least-busy
|
||||
router = Router(
|
||||
model_list=base_model_list,
|
||||
routing_strategy="simple-shuffle",
|
||||
model_list=model_list,
|
||||
routing_strategy="least-busy",
|
||||
)
|
||||
|
||||
original_kwargs = {"routing_strategy": "latency-based-routing", "other": "value"}
|
||||
request_kwargs = original_kwargs.copy()
|
||||
# Pre-populate latency cache for override strategy
|
||||
latency_cache = {
|
||||
"deployment-1": {"latency": [0.1]},
|
||||
"deployment-2": {"latency": [0.05]}, # Faster
|
||||
}
|
||||
router.cache.set_cache(key="test-model_map", value=latency_cache)
|
||||
|
||||
# Make request (lazy init will create logger)
|
||||
await router.async_get_available_deployment(
|
||||
model="gpt-3.5-turbo",
|
||||
request_kwargs=request_kwargs,
|
||||
# Override to latency-based
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="test-model",
|
||||
request_kwargs={"routing_strategy": "latency-based-routing"},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# After call, routing_strategy should be popped from request_kwargs
|
||||
assert "routing_strategy" not in request_kwargs
|
||||
# But other keys should remain
|
||||
assert "other" in request_kwargs
|
||||
# Should use latency-based routing (deployment-2), not least-busy
|
||||
assert deployment["model_info"]["id"] == "deployment-2"
|
||||
|
||||
def test_override_takes_precedence_over_global_sync(self):
|
||||
"""Test request override takes precedence in sync path."""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key1"},
|
||||
"model_info": {"id": "deployment-1"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "key2"},
|
||||
"model_info": {"id": "deployment-2"},
|
||||
},
|
||||
]
|
||||
|
||||
# Global: least-busy
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
routing_strategy="least-busy",
|
||||
)
|
||||
|
||||
# Pre-populate latency cache for override strategy
|
||||
latency_cache = {
|
||||
"deployment-1": {"latency": [0.1]},
|
||||
"deployment-2": {"latency": [0.05]}, # Faster
|
||||
}
|
||||
router.cache.set_cache(key="test-model_map", value=latency_cache)
|
||||
|
||||
# Override to latency-based
|
||||
deployment = router.get_available_deployment(
|
||||
model="test-model",
|
||||
request_kwargs={"routing_strategy": "latency-based-routing"},
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
)
|
||||
|
||||
# Should use latency-based routing (deployment-2), not least-busy
|
||||
assert deployment["model_info"]["id"] == "deployment-2"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue