ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: selectively apply routing strategy according to model name

Browse Source
This commit is contained in:
Yassin Kortam 2026-05-01 16:11:31 -07:00
parent 790d8bbe1a
commit 516b741de1
3 changed files with 1061 additions and 241 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

656
litellm/router.py
View File

@ -159,6 +159,7 @@ from litellm.types.router import (
RouterModelGroupAliasItem,
RouterRateLimitError,
RouterRateLimitErrorBasic,
RoutingGroup,
RoutingStrategy,
SearchToolTypedDict,
)
@ -308,6 +309,7 @@ class Router:
] = "simple-shuffle",
optional_pre_call_checks: Optional[OptionalPreCallChecks] = None,
routing_strategy_args: dict = {}, # just for latency-based
routing_groups: Optional[List[Union[RoutingGroup, dict]]] = None,
provider_budget_config: Optional[GenericBudgetConfigType] = None,
alerting_config: Optional[AlertingConfig] = None,
router_general_settings: Optional[
@ -347,8 +349,9 @@ class Router:
retry_after (int): Minimum time to wait before retrying a failed request. Defaults to 0.
allowed_fails (Optional[int]): Number of allowed fails before adding to cooldown. Defaults to None.
cooldown_time (float): Time to cooldown a deployment after failure in seconds. Defaults to 1.
routing_strategy (Literal["simple-shuffle", "least-busy", "usage-based-routing", "latency-based-routing", "cost-based-routing"]): Routing strategy. Defaults to "simple-shuffle".
routing_strategy_args (dict): Additional args for latency-based routing. Defaults to {}.
routing_strategy (Literal["simple-shuffle", "least-busy", "usage-based-routing", "latency-based-routing", "cost-based-routing"]): Routing strategy used for the implicit "default" group (any model not claimed by an entry in `routing_groups`). Defaults to "simple-shuffle".
routing_strategy_args (dict): Additional args for the default group's routing strategy (e.g. latency window). Defaults to {}.
routing_groups (Optional[List[RoutingGroup]]): Named subsets of `model_name`s that use a per-group routing strategy and args. Each model belongs to at most one explicit group; everything else lands in the implicit "default" group driven by `routing_strategy` / `routing_strategy_args`. Defaults to None.
alerting_config (AlertingConfig): Slack alerting configuration. Defaults to None.
provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None.
deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600.
@ -541,7 +544,10 @@ class Router:
self.stream_timeout = stream_timeout
self.retry_after = retry_after
self.routing_strategy = routing_strategy
self.routing_strategy = self._normalize_strategy(routing_strategy)
self._routing_groups_input: Optional[List[Union[RoutingGroup, dict]]] = (
routing_groups
)
## SETTING FALLBACKS ##
### validate if it's set + in correct format
@ -612,6 +618,7 @@ class Router:
routing_strategy=routing_strategy,
routing_strategy_args=routing_strategy_args,
)
self._init_routing_groups(self._routing_groups_input)
self.access_groups = None
## USAGE TRACKING ##
if isinstance(litellm._async_success_callback, list):
@ -806,84 +813,349 @@ class Router:
if self.cache.redis_cache is None:
self.cache.redis_cache = cache
# Maps a routing strategy string to the attribute on `self` that holds
# the default group's strategy selector for that strategy. (The selectors
# double as `CustomLogger` callbacks, hence the legacy `*_logger` attrs.)
_DEFAULT_SELECTOR_ATTR_BY_STRATEGY: Dict[str, str] = {
"least-busy": "leastbusy_logger",
"usage-based-routing": "lowesttpm_logger",
"usage-based-routing-v2": "lowesttpm_logger_v2",
"latency-based-routing": "lowestlatency_logger",
"cost-based-routing": "lowestcost_logger",
}
@staticmethod
def _normalize_strategy(
strategy: Union[RoutingStrategy, str, None]
) -> Optional[str]:
if strategy is None:
return None
if isinstance(strategy, RoutingStrategy):
return strategy.value
return strategy
def _validate_routing_strategy(
self, routing_strategy: Union[RoutingStrategy, str, None]
) -> None:
# See: https://github.com/BerriAI/litellm/issues/11330
valid_strategy_strings = ["simple-shuffle"] + [s.value for s in RoutingStrategy]
if routing_strategy is None:
return
is_valid_string = (
isinstance(routing_strategy, str)
and routing_strategy in valid_strategy_strings
)
is_valid_enum = isinstance(routing_strategy, RoutingStrategy)
if not is_valid_string and not is_valid_enum:
raise ValueError(
f"Invalid routing_strategy: '{routing_strategy}'. "
f"Valid options: {valid_strategy_strings}. "
f"Check 'router_settings.routing_strategy' in your config.yaml "
f"or the 'routing_strategy' parameter if using the Router SDK directly."
)
def _build_strategy_selector(
self,
strategy: Union[RoutingStrategy, str],
routing_strategy_args: dict,
register_callbacks: bool = True,
) -> Optional[Any]:
"""
Constructs a strategy selector for a given strategy.
Returns None for `simple-shuffle` (no selector needed) and unknown
strategies.
"""
selector: Optional[Any] = None
match self._normalize_strategy(strategy):
case RoutingStrategy.LEAST_BUSY.value:
selector = LeastBusyLoggingHandler(router_cache=self.cache)
if register_callbacks:
if isinstance(litellm.input_callback, list):
litellm.input_callback.append(selector) # type: ignore
else:
litellm.input_callback = [selector] # type: ignore
case RoutingStrategy.USAGE_BASED_ROUTING.value:
selector = LowestTPMLoggingHandler(
router_cache=self.cache,
routing_args=routing_strategy_args,
)
case RoutingStrategy.USAGE_BASED_ROUTING_V2.value:
selector = LowestTPMLoggingHandler_v2(
router_cache=self.cache,
routing_args=routing_strategy_args,
)
case RoutingStrategy.LATENCY_BASED.value:
selector = LowestLatencyLoggingHandler(
router_cache=self.cache,
routing_args=routing_strategy_args,
)
case RoutingStrategy.COST_BASED.value:
selector = LowestCostLoggingHandler(
router_cache=self.cache,
routing_args={},
)
if (
selector is not None
and register_callbacks
and isinstance(litellm.callbacks, list)
):
litellm.logging_callback_manager.add_litellm_callback(selector) # type: ignore
return selector
def _unregister_router_selectors(self, selectors: List[Any]) -> None:
"""
Drop router-owned strategy selectors from litellm's global callback
lists by identity. Used before re-init (`routing_strategy_init` /
`_init_routing_groups`) so repeated `update_settings` calls don't
accumulate dead selectors that keep receiving callback events.
"""
selector_ids = {id(s) for s in selectors if s is not None}
if not selector_ids:
return
if isinstance(litellm.callbacks, list):
litellm.callbacks = [
c for c in litellm.callbacks if id(c) not in selector_ids
]
if isinstance(litellm.input_callback, list):
litellm.input_callback = [
c for c in litellm.input_callback if id(c) not in selector_ids
]
def routing_strategy_init(
self, routing_strategy: Union[RoutingStrategy, str], routing_strategy_args: dict
):
verbose_router_logger.info(f"Routing strategy: {routing_strategy}")
self._validate_routing_strategy(routing_strategy)
# Validate routing_strategy value to fail fast with helpful error
# See: https://github.com/BerriAI/litellm/issues/11330
# Derive valid strategies from RoutingStrategy enum + "simple-shuffle" (default, not in enum)
valid_strategy_strings = ["simple-shuffle"] + [s.value for s in RoutingStrategy]
self._unregister_router_selectors(
[
getattr(self, attr, None)
for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()
]
)
if routing_strategy is not None:
is_valid_string = (
isinstance(routing_strategy, str)
and routing_strategy in valid_strategy_strings
)
is_valid_enum = isinstance(routing_strategy, RoutingStrategy)
if not is_valid_string and not is_valid_enum:
self.leastbusy_logger: Optional[LeastBusyLoggingHandler] = None
self.lowesttpm_logger: Optional[LowestTPMLoggingHandler] = None
self.lowesttpm_logger_v2: Optional[LowestTPMLoggingHandler_v2] = None
self.lowestlatency_logger: Optional[LowestLatencyLoggingHandler] = None
self.lowestcost_logger: Optional[LowestCostLoggingHandler] = None
selector = self._build_strategy_selector(
strategy=routing_strategy,
routing_strategy_args=routing_strategy_args,
)
attr = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(
self._normalize_strategy(routing_strategy) or ""
)
# TODO: legacy `self.<strategy>_logger` attributes are read directly by
# `get_settings()` and external callers. Fold the default group into
# `self._group_selectors["default"]` and drop these attribute writes —
# the dual storage is an antipattern preserved only for back-compat.
if attr is not None:
setattr(self, attr, selector)
def _init_routing_groups(
self,
groups_input: Optional[List[Union[RoutingGroup, dict]]],
) -> None:
"""
Validates and indexes `routing_groups`. Each `model_name` may belong to
at most one explicit group. Constructs per-group strategy selectors so
groups with different `routing_strategy_args` track independent state.
Models not claimed by any explicit group are served by the implicit
`"default"` group, whose selectors are the `self.<strategy>_logger`
attributes set up in `routing_strategy_init`.
"""
self._unregister_router_selectors(
[
sel
for selectors in getattr(self, "_group_selectors", {}).values()
for sel in selectors.values()
]
)
self._routing_groups: Dict[str, RoutingGroup] = {}
self._model_to_group: Dict[str, str] = {}
self._group_selectors: Dict[str, Dict[str, Any]] = {}
if not groups_input:
return
known_model_names = {
m.get("model_name") for m in (self.model_list or []) if m.get("model_name")
}
seen_group_names: set = set()
for raw in groups_input:
group = raw if isinstance(raw, RoutingGroup) else RoutingGroup(**raw)
if not group.group_name:
raise ValueError("routing_groups: group_name must be non-empty.")
if group.group_name == "default":
raise ValueError(
f"Invalid routing_strategy: '{routing_strategy}'. "
f"Valid options: {valid_strategy_strings}. "
f"Check 'router_settings.routing_strategy' in your config.yaml "
f"or the 'routing_strategy' parameter if using the Router SDK directly."
"routing_groups: 'default' is reserved for the implicit fallback group."
)
if group.group_name in seen_group_names:
raise ValueError(
f"routing_groups: group names must be unique, duplicate group_name '{group.group_name}'."
)
seen_group_names.add(group.group_name)
if (
routing_strategy == RoutingStrategy.LEAST_BUSY.value
or routing_strategy == RoutingStrategy.LEAST_BUSY
):
self.leastbusy_logger = LeastBusyLoggingHandler(router_cache=self.cache)
## add callback
if isinstance(litellm.input_callback, list):
litellm.input_callback.append(self.leastbusy_logger) # type: ignore
else:
litellm.input_callback = [self.leastbusy_logger] # type: ignore
if isinstance(litellm.callbacks, list):
litellm.logging_callback_manager.add_litellm_callback(self.leastbusy_logger) # type: ignore
elif (
routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING.value
or routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING
):
self.lowesttpm_logger = LowestTPMLoggingHandler(
router_cache=self.cache,
routing_args=routing_strategy_args,
self._validate_routing_strategy(group.routing_strategy)
for model_name in group.models:
if model_name in self._model_to_group:
raise ValueError(
f"routing_groups: model_name '{model_name}' appears in "
f"both '{self._model_to_group[model_name]}' and "
f"'{group.group_name}'. Each model may belong to at most one group."
)
if known_model_names and model_name not in known_model_names:
verbose_router_logger.warning(
"routing_groups: model_name '%s' (group '%s') is not in model_list; "
"the group entry will only take effect once a deployment with that "
"model_name is added.",
model_name,
group.group_name,
)
self._model_to_group[model_name] = group.group_name
self._routing_groups[group.group_name] = group
strategy_value = self._normalize_strategy(group.routing_strategy) or ""
group_selector = self._build_strategy_selector(
strategy=group.routing_strategy,
routing_strategy_args=group.routing_strategy_args or {},
)
if isinstance(litellm.callbacks, list):
litellm.logging_callback_manager.add_litellm_callback(self.lowesttpm_logger) # type: ignore
elif (
routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING_V2.value
or routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING_V2
):
self.lowesttpm_logger_v2 = LowestTPMLoggingHandler_v2(
router_cache=self.cache,
routing_args=routing_strategy_args,
self._group_selectors[group.group_name] = (
{strategy_value: group_selector} if group_selector is not None else {}
)
if isinstance(litellm.callbacks, list):
litellm.logging_callback_manager.add_litellm_callback(self.lowesttpm_logger_v2) # type: ignore
elif (
routing_strategy == RoutingStrategy.LATENCY_BASED.value
or routing_strategy == RoutingStrategy.LATENCY_BASED
):
self.lowestlatency_logger = LowestLatencyLoggingHandler(
router_cache=self.cache,
routing_args=routing_strategy_args,
)
if isinstance(litellm.callbacks, list):
litellm.logging_callback_manager.add_litellm_callback(self.lowestlatency_logger) # type: ignore
elif (
routing_strategy == RoutingStrategy.COST_BASED.value
or routing_strategy == RoutingStrategy.COST_BASED
):
self.lowestcost_logger = LowestCostLoggingHandler(
router_cache=self.cache,
routing_args={},
)
if isinstance(litellm.callbacks, list):
litellm.logging_callback_manager.add_litellm_callback(self.lowestcost_logger) # type: ignore
else:
pass
def _get_routing_context(self, model: str) -> Tuple[Optional[str], Optional[Any]]:
"""
Resolves the routing strategy and selector to use for the given model.
Every model belongs to exactly one group: an explicit entry from
`routing_groups`, or the implicit `"default"` group driven by the
router's top-level `routing_strategy` / `routing_strategy_args`.
`self.routing_strategy` may be either a string or a `RoutingStrategy`
enum member (the constructor accepts both), so it is normalized to a
string here. Downstream call sites and `_select_deployment_*` arms
compare against string literals.
"""
group_name = self._model_to_group.get(model)
if group_name is None:
strategy = self._normalize_strategy(self.routing_strategy)
attr = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "")
selector = getattr(self, attr, None) if attr is not None else None
return strategy, selector
group = self._routing_groups[group_name]
strategy = self._normalize_strategy(group.routing_strategy)
selector = self._group_selectors.get(group_name, {}).get(strategy or "")
return strategy, selector
async def _select_deployment_async(
self,
*,
strategy: Optional[str],
selector: Optional[Any],
model: str,
healthy_deployments: list,
messages: Optional[List[Dict[str, str]]],
input: Optional[Union[str, List]],
request_kwargs: Optional[Dict],
) -> Optional[Any]:
"""
Asks the strategy selector for a deployment. Caller handles
`simple-shuffle` separately (it does not flow through a selector).
Returns None for unknown strategies or when the selector is missing.
"""
if selector is None:
return None
match strategy:
case "least-busy":
return await selector.async_get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments,
)
case "usage-based-routing":
# `LowestTPMLoggingHandler` (v1) only exposes the sync
# `get_available_deployments`. Mirror the pre-routing-groups
# top-level fallback by calling it inline so groups using v1
# still work from async callers.
return selector.get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments,
messages=messages,
input=input,
)
case "usage-based-routing-v2" | "cost-based-routing":
return await selector.async_get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments,
messages=messages,
input=input,
)
case "latency-based-routing":
return await selector.async_get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments,
messages=messages,
input=input,
request_kwargs=request_kwargs,
)
case _:
return None
def _select_deployment_sync(
self,
*,
strategy: Optional[str],
selector: Optional[Any],
model: str,
healthy_deployments: list,
messages: Optional[List[Dict[str, str]]],
input: Optional[Union[str, List]],
request_kwargs: Optional[Dict],
) -> Optional[Any]:
"""
Sync sibling of `_select_deployment_async`. Caller handles
`simple-shuffle` separately.
"""
if selector is None:
return None
# `cost-based-routing` is intentionally omitted —
# `LowestCostLoggingHandler` only implements
# `async_get_available_deployments`
match strategy:
case "least-busy":
return selector.get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments,
)
case "usage-based-routing" | "usage-based-routing-v2":
return selector.get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments,
messages=messages,
input=input,
)
case "latency-based-routing":
return selector.get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments,
messages=messages,
input=input,
request_kwargs=request_kwargs,
)
case _:
return None
def initialize_assistants_endpoint(self):
## INITIALIZE PASS THROUGH ASSISTANTS ENDPOINT ##
@ -9016,8 +9288,13 @@ class Router:
if (
var == "routing_strategy_args"
and self.routing_strategy == "latency-based-routing"
and self.lowestlatency_logger is not None
):
_settings_to_return[var] = self.lowestlatency_logger.routing_args.json()
_settings_to_return["routing_groups"] = [
group.model_dump() for group in self._routing_groups.values()
]
return _settings_to_return
def update_settings(self, **kwargs):
@ -9028,6 +9305,7 @@ class Router:
_allowed_settings = [
"routing_strategy_args",
"routing_strategy",
"routing_groups",
"allowed_fails",
"cooldown_time",
"num_retries",
@ -9049,26 +9327,34 @@ class Router:
]
_existing_router_settings = self.get_settings()
rebuild_routing_groups = False
for var in kwargs:
if var in _allowed_settings:
if var in _int_settings:
_casted_value = int(kwargs[var])
setattr(self, var, _casted_value)
elif var == "routing_groups":
self._routing_groups_input = kwargs[var]
rebuild_routing_groups = True
else:
value = kwargs[var]
# only run routing strategy init if it has changed
if (
var == "routing_strategy"
and _existing_router_settings["routing_strategy"] != kwargs[var]
):
self.routing_strategy_init(
routing_strategy=kwargs[var],
routing_strategy_args=kwargs.get(
"routing_strategy_args", {}
),
)
setattr(self, var, kwargs[var])
if var == "routing_strategy":
value = self._normalize_strategy(value)
if _existing_router_settings["routing_strategy"] != value:
self.routing_strategy_init(
routing_strategy=value,
routing_strategy_args=kwargs.get(
"routing_strategy_args", {}
),
)
rebuild_routing_groups = True
setattr(self, var, value)
else:
verbose_router_logger.debug("Setting {} is not allowed".format(var))
if rebuild_routing_groups:
self._init_routing_groups(self._routing_groups_input)
verbose_router_logger.debug(f"Updated Router settings: {self.get_settings()}")
def _get_client(self, deployment, kwargs, client_type=None):
@ -9779,6 +10065,11 @@ class Router:
messages = pre_routing_hook_response.messages
#########################################################
# Resolve the strategy and logger AFTER the pre-routing hook, since
# the hook can replace `model` and routing-group lookup must key
# off the final model name.
strategy, strategy_selector = self._get_routing_context(model)
healthy_deployments = await self.async_get_healthy_deployments(
model=model,
request_kwargs=request_kwargs,
@ -9798,61 +10089,21 @@ class Router:
return healthy_deployments[0]
start_time = time.time()
if (
self.routing_strategy == "usage-based-routing-v2"
and self.lowesttpm_logger_v2 is not None
):
deployment = (
await self.lowesttpm_logger_v2.async_get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments, # type: ignore
messages=messages,
input=input,
)
)
elif (
self.routing_strategy == "cost-based-routing"
and self.lowestcost_logger is not None
):
deployment = (
await self.lowestcost_logger.async_get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments, # type: ignore
messages=messages,
input=input,
)
)
elif (
self.routing_strategy == "latency-based-routing"
and self.lowestlatency_logger is not None
):
deployment = (
await self.lowestlatency_logger.async_get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments, # type: ignore
messages=messages,
input=input,
request_kwargs=request_kwargs,
)
)
elif self.routing_strategy == "simple-shuffle":
if strategy == "simple-shuffle":
return simple_shuffle(
llm_router_instance=self,
healthy_deployments=healthy_deployments,
model=model,
)
elif (
self.routing_strategy == "least-busy"
and self.leastbusy_logger is not None
):
deployment = (
await self.leastbusy_logger.async_get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments, # type: ignore
)
)
else:
deployment = None
deployment = await self._select_deployment_async(
strategy=strategy,
selector=strategy_selector,
model=model,
healthy_deployments=healthy_deployments, # type: ignore
messages=messages,
input=input,
request_kwargs=request_kwargs,
)
if deployment is None:
exception = await async_raise_no_deployment_exception(
litellm_router_instance=self,
@ -9960,49 +10211,22 @@ class Router:
# 5. Apply load balancing strategy
start_time = time.perf_counter()
if (
self.routing_strategy == "usage-based-routing-v2"
and self.lowesttpm_logger_v2 is not None
):
deployment = (
await self.lowesttpm_logger_v2.async_get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
messages=messages,
input=input,
)
)
elif (
self.routing_strategy == "latency-based-routing"
and self.lowestlatency_logger is not None
):
deployment = (
await self.lowestlatency_logger.async_get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
messages=messages,
input=input,
request_kwargs=request_kwargs,
)
)
elif self.routing_strategy == "simple-shuffle":
strategy, strategy_selector = self._get_routing_context(model)
if strategy == "simple-shuffle":
return simple_shuffle(
llm_router_instance=self,
healthy_deployments=pass_through_deployments,
model=model,
)
elif (
self.routing_strategy == "least-busy"
and self.leastbusy_logger is not None
):
deployment = (
await self.leastbusy_logger.async_get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
)
)
else:
deployment = None
deployment = await self._select_deployment_async(
strategy=strategy,
selector=strategy_selector,
model=model,
healthy_deployments=pass_through_deployments, # type: ignore
messages=messages,
input=input,
request_kwargs=request_kwargs,
)
if deployment is None:
exception = await async_raise_no_deployment_exception(
@ -10191,11 +10415,8 @@ class Router:
cooldown_list=_cooldown_list,
)
if self.routing_strategy == "least-busy" and self.leastbusy_logger is not None:
deployment = self.leastbusy_logger.get_available_deployments(
model_group=model, healthy_deployments=healthy_deployments # type: ignore
)
elif self.routing_strategy == "simple-shuffle":
strategy, strategy_selector = self._get_routing_context(model)
if strategy == "simple-shuffle":
# if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm
############## Check 'weight' param set for weighted pick #################
return simple_shuffle(
@ -10203,37 +10424,15 @@ class Router:
healthy_deployments=healthy_deployments,
model=model,
)
elif (
self.routing_strategy == "latency-based-routing"
and self.lowestlatency_logger is not None
):
deployment = self.lowestlatency_logger.get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments, # type: ignore
request_kwargs=request_kwargs,
)
elif (
self.routing_strategy == "usage-based-routing"
and self.lowesttpm_logger is not None
):
deployment = self.lowesttpm_logger.get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments, # type: ignore
messages=messages,
input=input,
)
elif (
self.routing_strategy == "usage-based-routing-v2"
and self.lowesttpm_logger_v2 is not None
):
deployment = self.lowesttpm_logger_v2.get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments, # type: ignore
messages=messages,
input=input,
)
else:
deployment = None
deployment = self._select_deployment_sync(
strategy=strategy,
selector=strategy_selector,
model=model,
healthy_deployments=healthy_deployments, # type: ignore
messages=messages,
input=input,
request_kwargs=request_kwargs,
)
if deployment is None:
verbose_router_logger.info(
@ -10359,47 +10558,22 @@ class Router:
)
# 6. Apply load balancing strategy
if self.routing_strategy == "least-busy" and self.leastbusy_logger is not None:
deployment = self.leastbusy_logger.get_available_deployments(
model_group=model, healthy_deployments=pass_through_deployments # type: ignore
)
elif self.routing_strategy == "simple-shuffle":
strategy, strategy_selector = self._get_routing_context(model)
if strategy == "simple-shuffle":
return simple_shuffle(
llm_router_instance=self,
healthy_deployments=pass_through_deployments,
model=model,
)
elif (
self.routing_strategy == "latency-based-routing"
and self.lowestlatency_logger is not None
):
deployment = self.lowestlatency_logger.get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
request_kwargs=request_kwargs,
)
elif (
self.routing_strategy == "usage-based-routing"
and self.lowesttpm_logger is not None
):
deployment = self.lowesttpm_logger.get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
messages=messages,
input=input,
)
elif (
self.routing_strategy == "usage-based-routing-v2"
and self.lowesttpm_logger_v2 is not None
):
deployment = self.lowesttpm_logger_v2.get_available_deployments(
model_group=model,
healthy_deployments=pass_through_deployments, # type: ignore
messages=messages,
input=input,
)
else:
deployment = None
deployment = self._select_deployment_sync(
strategy=strategy,
selector=strategy_selector,
model=model,
healthy_deployments=pass_through_deployments, # type: ignore
messages=messages,
input=input,
request_kwargs=request_kwargs,
)
if deployment is None:
verbose_router_logger.info(

15
litellm/types/router.py
View File

@ -38,6 +38,19 @@ class ModelConfig(BaseModel):
model_config = ConfigDict(protected_namespaces=())
class RoutingGroup(BaseModel):
"""
A group of models that share a routing strategy.
"""
group_name: str
models: List[str]
routing_strategy: str
routing_strategy_args: Optional[dict] = None
model_config = ConfigDict(protected_namespaces=())
class RouterConfig(BaseModel):
model_list: List[ModelConfig]
@ -65,6 +78,7 @@ class RouterConfig(BaseModel):
"usage-based-routing",
"latency-based-routing",
] = "simple-shuffle"
routing_groups: Optional[List[RoutingGroup]] = None
model_config = ConfigDict(protected_namespaces=())
@ -76,6 +90,7 @@ class UpdateRouterConfig(BaseModel):
routing_strategy_args: Optional[dict] = None
routing_strategy: Optional[str] = None
routing_groups: Optional[List[RoutingGroup]] = None
model_group_retry_policy: Optional[dict] = None
model_group_affinity_config: Optional[Dict[str, List[str]]] = None
allowed_fails: Optional[int] = None

631
tests/test_litellm/router_strategy/test_router_routing_groups.py Normal file
View File

@ -0,0 +1,631 @@
"""
Tests for `routing_groups` — assigns named subsets of `model_name`s to a
per-group routing strategy. Models not claimed by an explicit group fall into
the implicit `"default"` group driven by the router's top-level
`routing_strategy` / `routing_strategy_args`.
"""
import os
import sys
from unittest.mock import patch
import pytest
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm import Router
from litellm.types.router import RoutingGroup, RoutingStrategy
def _model_list():
return [
{
"model_name": "filtered-model",
"litellm_params": {
"model": "openai/gpt-4o",
"api_key": "sk-test-1",
"api_base": "https://example.invalid",
},
"model_info": {"id": "deploy-1"},
},
{
"model_name": "filtered-model",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "sk-test-2",
"api_base": "https://example.invalid",
},
"model_info": {"id": "deploy-2"},
},
{
"model_name": "other-model",
"litellm_params": {
"model": "openai/gpt-4o",
"api_key": "sk-test-3",
"api_base": "https://example.invalid",
},
"model_info": {"id": "deploy-3"},
},
]
def _build_router(routing_strategy="simple-shuffle", routing_groups=None):
return Router(
model_list=_model_list(),
routing_strategy=routing_strategy,
routing_groups=routing_groups,
)
def test_no_groups_uses_top_level_strategy_for_all_models():
router = _build_router(routing_strategy="latency-based-routing")
assert router._get_routing_context("filtered-model")[0] == "latency-based-routing"
assert router._get_routing_context("other-model")[0] == "latency-based-routing"
def test_explicit_group_overrides_top_level():
router = _build_router(
routing_strategy="simple-shuffle",
routing_groups=[
{
"group_name": "fast",
"models": ["filtered-model"],
"routing_strategy": "latency-based-routing",
}
],
)
strategy, selector = router._get_routing_context("filtered-model")
assert strategy == "latency-based-routing"
assert selector is not None
# other-model isn't in any explicit group, so it lands in the default fallback
strategy_other, _ = router._get_routing_context("other-model")
assert strategy_other == "simple-shuffle"
def test_default_group_is_simple_shuffle_when_top_level_strategy_unset():
router = _build_router()
strategy, selector = router._get_routing_context("other-model")
assert strategy == "simple-shuffle"
# simple-shuffle does not use a selector
assert selector is None
def test_two_groups_same_strategy_have_independent_selectors():
router = _build_router(
routing_strategy="simple-shuffle",
routing_groups=[
{
"group_name": "group_a",
"models": ["filtered-model"],
"routing_strategy": "latency-based-routing",
"routing_strategy_args": {"ttl": 60},
},
{
"group_name": "group_b",
"models": ["other-model"],
"routing_strategy": "latency-based-routing",
"routing_strategy_args": {"ttl": 600},
},
],
)
a_selector = router._group_selectors["group_a"]["latency-based-routing"]
b_selector = router._group_selectors["group_b"]["latency-based-routing"]
assert a_selector is not None and b_selector is not None
assert id(a_selector) != id(b_selector)
def test_overlapping_models_across_groups_raises():
with pytest.raises(ValueError, match="appears in"):
_build_router(
routing_groups=[
{
"group_name": "g1",
"models": ["filtered-model"],
"routing_strategy": "latency-based-routing",
},
{
"group_name": "g2",
"models": ["filtered-model"],
"routing_strategy": "least-busy",
},
],
)
def test_reserved_default_group_name_raises():
with pytest.raises(ValueError, match="reserved"):
_build_router(
routing_groups=[
{
"group_name": "default",
"models": ["filtered-model"],
"routing_strategy": "latency-based-routing",
}
],
)
def test_invalid_strategy_in_group_raises():
with pytest.raises(ValueError, match="Invalid routing_strategy"):
_build_router(
routing_groups=[
{
"group_name": "g1",
"models": ["filtered-model"],
"routing_strategy": "not-a-real-strategy",
}
],
)
def test_unknown_model_in_group_warns_but_does_not_raise(caplog):
import logging
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
router = _build_router(
routing_groups=[
{
"group_name": "g1",
"models": ["model-not-in-list"],
"routing_strategy": "latency-based-routing",
}
],
)
assert "model-not-in-list" in caplog.text
assert router._model_to_group.get("model-not-in-list") == "g1"
def test_duplicate_group_name_raises():
with pytest.raises(ValueError, match="duplicate"):
_build_router(
routing_groups=[
{
"group_name": "g1",
"models": ["filtered-model"],
"routing_strategy": "latency-based-routing",
},
{
"group_name": "g1",
"models": ["other-model"],
"routing_strategy": "least-busy",
},
],
)
def test_routing_group_object_input_accepted():
router = _build_router(
routing_groups=[
RoutingGroup(
group_name="g1",
models=["filtered-model"],
routing_strategy="latency-based-routing",
)
],
)
assert router._model_to_group["filtered-model"] == "g1"
@pytest.mark.asyncio
async def test_async_dispatch_uses_group_strategy_for_grouped_model():
router = _build_router(
routing_strategy="simple-shuffle",
routing_groups=[
{
"group_name": "fast",
"models": ["filtered-model"],
"routing_strategy": "latency-based-routing",
}
],
)
group_selector = router._group_selectors["fast"]["latency-based-routing"]
with (
patch.object(
group_selector,
"async_get_available_deployments",
wraps=group_selector.async_get_available_deployments,
) as latency_spy,
patch(
"litellm.router.simple_shuffle", wraps=litellm.router.simple_shuffle
) as shuffle_spy,
):
await router.async_get_available_deployment(
model="filtered-model", request_kwargs={}
)
assert (
latency_spy.called
), "group's latency selector should run for grouped model"
assert not shuffle_spy.called
@pytest.mark.asyncio
async def test_async_dispatch_falls_back_to_default_for_ungrouped_models():
router = _build_router(
routing_strategy="simple-shuffle",
routing_groups=[
{
"group_name": "fast",
"models": ["filtered-model"],
"routing_strategy": "latency-based-routing",
}
],
)
group_selector = router._group_selectors["fast"]["latency-based-routing"]
with (
patch.object(
group_selector,
"async_get_available_deployments",
wraps=group_selector.async_get_available_deployments,
) as latency_spy,
patch(
"litellm.router.simple_shuffle", wraps=litellm.router.simple_shuffle
) as shuffle_spy,
):
await router.async_get_available_deployment(
model="other-model", request_kwargs={}
)
assert shuffle_spy.called, "default group's simple-shuffle should run"
assert not latency_spy.called
def test_update_settings_round_trip_routing_groups():
router = _build_router()
assert router._model_to_group == {}
router.update_settings(
routing_groups=[
{
"group_name": "fast",
"models": ["filtered-model"],
"routing_strategy": "latency-based-routing",
}
]
)
assert router._model_to_group == {"filtered-model": "fast"}
assert router._get_routing_context("filtered-model")[0] == "latency-based-routing"
settings = router.get_settings()
assert any(g["group_name"] == "fast" for g in settings["routing_groups"])
def test_default_group_accepts_routing_strategy_enum_for_top_level_strategy():
"""
The Router constructor accepts both string and RoutingStrategy enum forms.
The default group must resolve a real selector for either input shape, or
every ungrouped model raises NoDeploymentAvailable.
"""
router = _build_router(routing_strategy=RoutingStrategy.LATENCY_BASED)
strategy, selector = router._get_routing_context("other-model")
assert strategy == "latency-based-routing"
assert selector is not None
assert router.lowestlatency_logger is selector
@pytest.mark.asyncio
async def test_async_dispatch_uses_default_selector_when_constructed_with_enum():
router = _build_router(routing_strategy=RoutingStrategy.LATENCY_BASED)
default_selector = router.lowestlatency_logger
assert default_selector is not None
# filtered-model has 2 deployments, so the single-deployment short-circuit
# in async_get_available_deployment doesn't kick in and we actually hit
# the latency selector.
with (
patch.object(
default_selector,
"async_get_available_deployments",
wraps=default_selector.async_get_available_deployments,
) as latency_spy,
patch(
"litellm.router.simple_shuffle", wraps=litellm.router.simple_shuffle
) as shuffle_spy,
):
await router.async_get_available_deployment(
model="filtered-model", request_kwargs={}
)
assert (
latency_spy.called
), "default group's latency selector must run when constructed with the enum"
assert not shuffle_spy.called
def test_update_settings_does_not_leak_strategy_callbacks(monkeypatch):
"""
Repeated `update_settings` calls must not accumulate stale selectors in
`litellm.callbacks` / `litellm.input_callback`. Each rebuild owns the
previous generation and is responsible for unregistering it by identity
so the global lists don't grow without bound.
"""
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm, "input_callback", [])
router = _build_router(routing_strategy="latency-based-routing")
initial_default_selector = router.lowestlatency_logger
assert initial_default_selector is not None
assert any(
c is initial_default_selector for c in litellm.callbacks
), "fresh router should register its default selector"
# Flip strategies back and forth — each transition triggers
# routing_strategy_init and must replace, not accumulate.
for _ in range(3):
router.update_settings(routing_strategy="usage-based-routing")
router.update_settings(routing_strategy="latency-based-routing")
assert all(
c is not initial_default_selector for c in litellm.callbacks
), "old default selector instance leaked into litellm.callbacks"
# The set of router-owned selectors should be bounded — at most one per
# strategy class. Without the cleanup fix this grows by one per toggle.
router_owned_classes = (
"LeastBusyLoggingHandler",
"LowestTPMLoggingHandler",
"LowestTPMLoggingHandler_v2",
"LowestLatencyLoggingHandler",
"LowestCostLoggingHandler",
)
router_owned = [
c for c in litellm.callbacks if type(c).__name__ in router_owned_classes
]
assert len(router_owned) <= len(
router_owned_classes
), f"router selectors leaked: {[type(c).__name__ for c in router_owned]}"
# Group selectors: capture the v1 instance, swap, confirm v1 is gone by id.
router.update_settings(
routing_groups=[
{
"group_name": "fast",
"models": ["filtered-model"],
"routing_strategy": "latency-based-routing",
}
]
)
group_selector_v1 = router._group_selectors["fast"]["latency-based-routing"]
assert group_selector_v1 is not None
router.update_settings(
routing_groups=[
{
"group_name": "fast",
"models": ["filtered-model"],
"routing_strategy": "latency-based-routing",
"routing_strategy_args": {"ttl": 600},
}
]
)
assert all(
c is not group_selector_v1 for c in litellm.callbacks
), "old group selector instance leaked into litellm.callbacks"
group_selector_v2 = router._group_selectors["fast"]["latency-based-routing"]
assert group_selector_v2 is not group_selector_v1
def test_update_settings_unregisters_group_selectors_when_groups_removed(monkeypatch):
"""
Setting routing_groups to an empty list (or omitting all groups) must
unregister every previously-owned group selector.
"""
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm, "input_callback", [])
router = _build_router(
routing_groups=[
{
"group_name": "fast",
"models": ["filtered-model"],
"routing_strategy": "least-busy",
}
]
)
group_selector = router._group_selectors["fast"]["least-busy"]
assert group_selector in litellm.callbacks
assert group_selector in litellm.input_callback
router.update_settings(routing_groups=[])
assert all(c is not group_selector for c in litellm.callbacks)
assert all(c is not group_selector for c in litellm.input_callback)
assert router._group_selectors == {}
# ---------------------------------------------------------------------------
# Direct helper coverage
# ---------------------------------------------------------------------------
def test_normalize_strategy_handles_string_enum_and_none():
assert Router._normalize_strategy(None) is None
assert Router._normalize_strategy("simple-shuffle") == "simple-shuffle"
assert (
Router._normalize_strategy(RoutingStrategy.LATENCY_BASED)
== RoutingStrategy.LATENCY_BASED.value
)
def test_validate_routing_strategy_accepts_valid_and_rejects_invalid():
router = _build_router()
# Valid: string, enum, None
router._validate_routing_strategy("simple-shuffle")
router._validate_routing_strategy(RoutingStrategy.LATENCY_BASED)
router._validate_routing_strategy(None)
with pytest.raises(ValueError, match="Invalid routing_strategy"):
router._validate_routing_strategy("not-a-real-strategy")
def test_build_strategy_selector_returns_none_for_simple_shuffle(monkeypatch):
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm, "input_callback", [])
router = _build_router()
assert (
router._build_strategy_selector(
strategy="simple-shuffle",
routing_strategy_args={},
register_callbacks=False,
)
is None
)
def test_build_strategy_selector_constructs_for_known_strategies(monkeypatch):
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm, "input_callback", [])
router = _build_router()
selector = router._build_strategy_selector(
strategy="latency-based-routing",
routing_strategy_args={"ttl": 30},
register_callbacks=False,
)
assert selector is not None
def test_unregister_router_selectors_removes_by_identity(monkeypatch):
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm, "input_callback", [])
router = _build_router()
selector = router._build_strategy_selector(
strategy="least-busy",
routing_strategy_args={},
register_callbacks=True,
)
assert selector in litellm.callbacks
assert selector in litellm.input_callback
router._unregister_router_selectors([selector])
assert all(c is not selector for c in litellm.callbacks)
assert all(c is not selector for c in litellm.input_callback)
def test_init_routing_groups_with_none_clears_state():
router = _build_router(
routing_groups=[
{
"group_name": "fast",
"models": ["filtered-model"],
"routing_strategy": "least-busy",
}
]
)
assert router._routing_groups
router._init_routing_groups(None)
assert router._routing_groups == {}
assert router._model_to_group == {}
assert router._group_selectors == {}
@pytest.mark.asyncio
async def test_select_deployment_async_returns_none_without_selector():
router = _build_router()
result = await router._select_deployment_async(
strategy="simple-shuffle",
selector=None,
model="filtered-model",
healthy_deployments=[],
messages=None,
input=None,
request_kwargs=None,
)
assert result is None
def test_select_deployment_sync_returns_none_without_selector():
router = _build_router()
result = router._select_deployment_sync(
strategy="simple-shuffle",
selector=None,
model="filtered-model",
healthy_deployments=[],
messages=None,
input=None,
request_kwargs=None,
)
assert result is None
def test_select_deployment_sync_does_not_dispatch_cost_based_routing():
"""
`LowestCostLoggingHandler` only implements `async_get_available_deployments`,
so the sync dispatch must fall through to None for `cost-based-routing`
instead of raising AttributeError.
"""
router = _build_router(
routing_strategy="simple-shuffle",
routing_groups=[
{
"group_name": "cheap",
"models": ["filtered-model"],
"routing_strategy": "cost-based-routing",
}
],
)
cost_selector = router._group_selectors["cheap"]["cost-based-routing"]
assert not hasattr(cost_selector, "get_available_deployments"), (
"test premise broken: LowestCostLoggingHandler unexpectedly grew a "
"sync method — revisit the sync dispatch arm."
)
result = router._select_deployment_sync(
strategy="cost-based-routing",
selector=cost_selector,
model="filtered-model",
healthy_deployments=[],
messages=None,
input=None,
request_kwargs=None,
)
assert result is None
@pytest.mark.asyncio
async def test_async_dispatch_falls_back_to_sync_for_usage_based_routing_v1():
"""
`LowestTPMLoggingHandler` (v1) only implements the sync
`get_available_deployments`. The async dispatch must call the sync method
on the v1 selector instead of awaiting a non-existent async one.
"""
router = _build_router(
routing_strategy="simple-shuffle",
routing_groups=[
{
"group_name": "v1",
"models": ["filtered-model"],
"routing_strategy": "usage-based-routing",
}
],
)
v1_selector = router._group_selectors["v1"]["usage-based-routing"]
assert not hasattr(v1_selector, "async_get_available_deployments"), (
"test premise broken: LowestTPMLoggingHandler v1 unexpectedly grew an "
"async method — revisit the async dispatch arm."
)
with patch.object(
v1_selector,
"get_available_deployments",
wraps=v1_selector.get_available_deployments,
) as v1_spy:
await router._select_deployment_async(
strategy="usage-based-routing",
selector=v1_selector,
model="filtered-model",
healthy_deployments=[
{
"model_name": "filtered-model",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "deploy-1"},
}
],
messages=[{"role": "user", "content": "hi"}],
input=None,
request_kwargs={},
)
assert v1_spy.called, "async dispatch must route v1 strategy through sync method"