mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(auto_router): build the semantic route layer off the event loop
AutoRouter's cold-start route layer construction (SemanticRouter with auto_sync="local") ran directly on the event loop, doing at least one synchronous embedding HTTP call inline behind a bare "if routelayer is None" check with no lock, so it blocked the whole worker and let concurrent cold-start requests each build a duplicate layer. ComplexityRouter already solved this identically for its own semantic keyword matching (_ensure_semantic_routelayer: a lock plus asyncio.to_thread). Give AutoRouter the same treatment: extract the build into _build_routelayer and gate it behind _ensure_routelayer's double-checked async lock. Fixes #33204.
This commit is contained in:
parent
0cb759772c
commit
5fc769c0a1
2 changed files with 95 additions and 13 deletions
|
|
@ -2,6 +2,7 @@
|
|||
Auto-Routing Strategy that works with a Semantic Router Config
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional
|
||||
|
||||
|
|
@ -75,6 +76,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: SemanticRouter | None = None
|
||||
self._routelayer_lock = asyncio.Lock()
|
||||
self.default_model = default_model
|
||||
self.embedding_model: str = embedding_model
|
||||
self.max_input_chars: int = max_input_chars
|
||||
|
|
@ -115,6 +117,42 @@ class AutoRouter(CustomLogger):
|
|||
)
|
||||
return auto_router_routes
|
||||
|
||||
def _build_routelayer(self) -> "SemanticRouter":
|
||||
"""Build (once) the SemanticRouter for this alias's static route config.
|
||||
|
||||
`auto_sync="local"` embeds every route's utterances against the encoder, so
|
||||
this does a synchronous embedding call and must never run directly on the
|
||||
event loop; see `_ensure_routelayer`.
|
||||
"""
|
||||
if self.routelayer is not None:
|
||||
return self.routelayer
|
||||
|
||||
from semantic_router.routers import SemanticRouter
|
||||
|
||||
routelayer: Final = SemanticRouter(
|
||||
routes=self.loaded_routes,
|
||||
encoder=self.encoder,
|
||||
auto_sync=self.auto_sync_value,
|
||||
)
|
||||
self.routelayer = routelayer
|
||||
return routelayer
|
||||
|
||||
async def _ensure_routelayer(self) -> "SemanticRouter":
|
||||
"""Return the cached route layer, building it once under a lock if needed.
|
||||
|
||||
The build embeds the static route utterances via the encoder's synchronous
|
||||
path, so it runs in a worker thread to avoid blocking the event loop. A
|
||||
double-checked asyncio lock ensures concurrent cold-start requests build it
|
||||
exactly once rather than each firing duplicate embedding calls.
|
||||
"""
|
||||
if self.routelayer is not None:
|
||||
return self.routelayer
|
||||
async with self._routelayer_lock:
|
||||
routelayer = self.routelayer
|
||||
if routelayer is None:
|
||||
routelayer = await asyncio.to_thread(self._build_routelayer)
|
||||
return routelayer
|
||||
|
||||
@staticmethod
|
||||
def _extract_text_from_messages(messages: list[dict[str, Any]]) -> str:
|
||||
"""
|
||||
|
|
@ -151,8 +189,6 @@ class AutoRouter(CustomLogger):
|
|||
|
||||
Used for the litellm auto-router to modify the request before the routing decision is made.
|
||||
"""
|
||||
from semantic_router.routers import SemanticRouter
|
||||
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
|
||||
|
|
@ -164,17 +200,7 @@ class AutoRouter(CustomLogger):
|
|||
if resolved_messages is None:
|
||||
return None
|
||||
|
||||
routelayer = self.routelayer
|
||||
if routelayer is None:
|
||||
#######################
|
||||
# Create the route layer
|
||||
#######################
|
||||
routelayer = SemanticRouter(
|
||||
routes=self.loaded_routes,
|
||||
encoder=self.encoder,
|
||||
auto_sync=self.auto_sync_value,
|
||||
)
|
||||
self.routelayer = routelayer
|
||||
routelayer = await self._ensure_routelayer()
|
||||
|
||||
message_content: Final = self._extract_text_from_messages(resolved_messages)
|
||||
route_name: Final = await self._matched_route_name(routelayer, message_content, request_kwargs)
|
||||
|
|
|
|||
|
|
@ -604,3 +604,59 @@ class TestAutoRouterAttributesItsEmbeddingSpend:
|
|||
assert router.aembedding_kwargs["proxy_server_request"] == {
|
||||
"body": {"model": "text-embedding-3-small", "input": ["fix this stack trace"]}
|
||||
}
|
||||
|
||||
|
||||
class TestAutoRouterColdStartDoesNotBlockTheEventLoop:
|
||||
"""The first request through a fresh alias builds the route layer off the event loop thread,
|
||||
and concurrent first requests build it exactly once."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_build_the_routelayer_on_a_worker_thread_not_the_event_loop_thread(self):
|
||||
import threading
|
||||
|
||||
auto_router: Final = _auto_router(None, litellm_router_instance=StubEmbeddingRouter())
|
||||
event_loop_thread: Final = threading.get_ident()
|
||||
build_thread: list[int] = []
|
||||
original_build = auto_router._build_routelayer
|
||||
|
||||
def _tracking_build() -> Any:
|
||||
build_thread.append(threading.get_ident())
|
||||
return original_build()
|
||||
|
||||
auto_router._build_routelayer = _tracking_build # type: ignore[method-assign]
|
||||
|
||||
result: Final = await auto_router.async_pre_routing_hook(
|
||||
model="my-auto-router",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "fix this stack trace"}],
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert len(build_thread) == 1
|
||||
assert build_thread[0] != event_loop_thread
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_build_the_routelayer_exactly_once_under_concurrent_cold_start_requests(self):
|
||||
auto_router: Final = _auto_router(None, litellm_router_instance=StubEmbeddingRouter())
|
||||
build_calls: Final[list[int]] = []
|
||||
original_build = auto_router._build_routelayer
|
||||
|
||||
def _counting_build() -> Any:
|
||||
build_calls.append(1)
|
||||
return original_build()
|
||||
|
||||
auto_router._build_routelayer = _counting_build # type: ignore[method-assign]
|
||||
|
||||
results: Final = await asyncio.gather(
|
||||
*(
|
||||
auto_router.async_pre_routing_hook(
|
||||
model="my-auto-router",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "fix this stack trace"}],
|
||||
)
|
||||
for _ in range(10)
|
||||
)
|
||||
)
|
||||
|
||||
assert all(result is not None for result in results)
|
||||
assert len(build_calls) == 1
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue