diff --git a/litellm/router.py b/litellm/router.py index 8e29f8cfc1..8458c880e5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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._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._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( diff --git a/litellm/types/router.py b/litellm/types/router.py index 2bc8cee351..926815ba31 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -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 diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py new file mode 100644 index 0000000000..7368c72836 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -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"