diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index c01e2f10d2c..cd1a00b4fab 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -2,6 +2,7 @@ Auto-Routing Strategy that works with a Semantic Router Config """ +import asyncio from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from litellm._logging import verbose_router_logger @@ -48,6 +49,7 @@ class AutoRouter(CustomLogger): self.auto_sync_value = self.DEFAULT_AUTO_SYNC_VALUE self.loaded_routes: List[Route] = self._load_semantic_routing_routes() self.routelayer: Optional[SemanticRouter] = None + self._routelayer_lock = asyncio.Lock() self.default_model = default_model self.embedding_model: str = embedding_model self.litellm_router_instance: "Router" = litellm_router_instance @@ -132,21 +134,24 @@ class AutoRouter(CustomLogger): routelayer = self.routelayer if routelayer is None: - ####################### - # Create the route layer - ####################### - routelayer = SemanticRouter( - routes=self.loaded_routes, - encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.litellm_router_instance, - model_name=self.embedding_model, - ), - auto_sync=self.auto_sync_value, - ) - self.routelayer = routelayer + async with self._routelayer_lock: + routelayer = self.routelayer + if routelayer is None: + routelayer = await asyncio.to_thread( + SemanticRouter, + routes=self.loaded_routes, + encoder=LiteLLMRouterEncoder( + litellm_router_instance=self.litellm_router_instance, + model_name=self.embedding_model, + ), + auto_sync=self.auto_sync_value, + ) + self.routelayer = routelayer message_content = self._extract_text_from_messages(messages) - route_choice: Optional[Union[RouteChoice, List[RouteChoice]]] = routelayer(text=message_content) + route_choice: Optional[Union[RouteChoice, List[RouteChoice]]] = await asyncio.to_thread( + routelayer, text=message_content + ) verbose_router_logger.debug(f"route_choice: {route_choice}") if isinstance(route_choice, RouteChoice): model = route_choice.name or self.default_model diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/test_litellm/router_strategy/test_auto_router.py index cb46a4ae553..3a5989b6790 100644 --- a/tests/test_litellm/router_strategy/test_auto_router.py +++ b/tests/test_litellm/router_strategy/test_auto_router.py @@ -316,3 +316,87 @@ class TestAutoRouter: # Assert assert result is None + + +@pytest.fixture +def async_auto_router(mock_router_instance): + class RouteChoice: + def __init__(self, name): + self.name = name + + mock_semantic_router_class = MagicMock() + mock_semantic_router_class.from_json.return_value.routes = ["route1"] + mock_routelayer = MagicMock(return_value=RouteChoice(name="test-model")) + mock_semantic_router_class.return_value = mock_routelayer + modules = { + "semantic_router": MagicMock(), + "semantic_router.routers": MagicMock(SemanticRouter=mock_semantic_router_class), + "semantic_router.schema": MagicMock(RouteChoice=RouteChoice), + "litellm.router_strategy.auto_router.litellm_encoder": MagicMock(LiteLLMRouterEncoder=MagicMock()), + } + + with patch.dict("sys.modules", modules): + auto_router = AutoRouter( + model_name="test-auto-router", + auto_router_config_path="test/path/router.json", + default_model="gpt-4o-mini", + embedding_model="text-embedding-model", + litellm_router_instance=mock_router_instance, + ) + yield auto_router, mock_semantic_router_class, mock_routelayer + + +class TestAutoRouterAsyncExecution: + @pytest.mark.asyncio + async def test_offloads_construction_and_route_call(self, async_auto_router): + auto_router, mock_semantic_router_class, mock_routelayer = async_auto_router + + async def run_sync(function, *args, **kwargs): + return function(*args, **kwargs) + + with patch( + "litellm.router_strategy.auto_router.auto_router.asyncio.to_thread", + new=AsyncMock(side_effect=run_sync), + ) as mock_to_thread: + result = await auto_router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "test message"}], + ) + + assert result is not None + assert result.model == "test-model" + assert mock_to_thread.await_count == 2 + assert mock_to_thread.await_args_list[0].args[0] is mock_semantic_router_class + assert mock_to_thread.await_args_list[1].args == (mock_routelayer,) + assert mock_to_thread.await_args_list[1].kwargs == {"text": "test message"} + + @pytest.mark.asyncio + async def test_concurrent_cold_start_constructs_routelayer_once(self, async_auto_router): + auto_router, mock_semantic_router_class, mock_routelayer = async_auto_router + + async def run_sync(function, *args, **kwargs): + if function is mock_semantic_router_class: + await asyncio.sleep(0) + return function(*args, **kwargs) + + with patch( + "litellm.router_strategy.auto_router.auto_router.asyncio.to_thread", + new=AsyncMock(side_effect=run_sync), + ): + results = await asyncio.gather( + auto_router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "first message"}], + ), + auto_router.async_pre_routing_hook( + model="test-model", + request_kwargs={}, + messages=[{"role": "user", "content": "second message"}], + ), + ) + + assert all(result is not None for result in results) + mock_semantic_router_class.assert_called_once() + assert mock_routelayer.call_count == 2