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:
moe-berri 2026-09-05 15:03:05 -07:00
parent 0cb759772c
commit 5fc769c0a1
2 changed files with 95 additions and 13 deletions

View file

@ -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)

View file

@ -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