diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index f5decbaa670..2dae9393510 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -149,8 +149,11 @@ class ComplexityRouter(CustomLogger): self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS # Lazily built on first semantic request and cached for reuse (route - # embeddings are static, only the prompt is embedded per request). + # embeddings are static, only the prompt is embedded per request). The lock + # serializes the one-time build so concurrent cold-start requests don't each + # construct the index and fire duplicate embedding calls. self._semantic_routelayer: Optional[SemanticRouter] = None + self._semantic_routelayer_lock = asyncio.Lock() # Pre-compile regex patterns for efficiency # Use non-greedy .*? to prevent ReDoS on pathological inputs @@ -486,6 +489,22 @@ class ComplexityRouter(CustomLogger): self._semantic_routelayer = routelayer return routelayer + async def _ensure_semantic_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._semantic_routelayer is not None: + return self._semantic_routelayer + async with self._semantic_routelayer_lock: + routelayer = self._semantic_routelayer + if routelayer is None: + routelayer = await asyncio.to_thread(self._get_or_create_semantic_routelayer) + return routelayer + async def _semantic_tier_override(self, user_message: str, request_kwargs: Dict) -> Optional[ComplexityTier]: """Match the prompt against keyword_tier_rules by embedding similarity. @@ -503,11 +522,7 @@ class ComplexityRouter(CustomLogger): LiteLLMRouterEncoder, ) - # Building the SemanticRouter embeds the (static) route utterances via the - # encoder's *synchronous* path; run it in a worker thread so that one-time, - # per-router-instance provider I/O never blocks the async event loop and stalls - # other requests. Once cached, subsequent calls return immediately (no I/O). - routelayer = await asyncio.to_thread(self._get_or_create_semantic_routelayer) + routelayer = await self._ensure_semantic_routelayer() encoder = cast(LiteLLMRouterEncoder, routelayer.encoder) # cast-ok: always the encoder we built above # Strip the parent request's budget reservation before forwarding: the reservation # belongs to the routed completion this embedding is helping select, not to the diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index f4da1018f3e..eda73f7e7b6 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -4,6 +4,7 @@ Tests for the ComplexityRouter. Tests the rule-based complexity scoring and tier assignment logic. """ +import asyncio import os import sys from typing import Dict, List @@ -1477,6 +1478,11 @@ class FakeEmbeddingRouter: def __init__(self): self.async_embedding_calls: List[List[str]] = [] self.async_embedding_kwargs: List[Dict] = [] + # Every embedded batch (sync route-index build AND async query), so tests can count + # builds independently of which embedding path the library happens to use. + self.embedded_batches: List[List[str]] = [] + # Thread ids of the synchronous (route-index build) embedding calls, so a test can + # assert the build is offloaded off the event-loop thread. self.sync_embedding_thread_ids: List[int] = [] def _vectors(self, docs: List[str]) -> List[List[float]]: @@ -1490,19 +1496,24 @@ class FakeEmbeddingRouter: return text if isinstance(text, list) else [text] def embedding(self, input, model, **kwargs): - # Record the thread this synchronous (route-index build) call ran on, so tests can - # assert it is offloaded off the event-loop thread. import threading + docs = self._as_list(input) + self.embedded_batches.append(docs) self.sync_embedding_thread_ids.append(threading.get_ident()) - return _make_embedding_response(self._vectors(self._as_list(input))) + return _make_embedding_response(self._vectors(docs)) async def aembedding(self, input, model, **kwargs): docs = self._as_list(input) + self.embedded_batches.append(docs) self.async_embedding_calls.append(docs) self.async_embedding_kwargs.append(kwargs) return _make_embedding_response(self._vectors(docs)) + def utterance_embedding_count(self, utterance: str) -> int: + """How many times the given route utterance was embedded == number of route-index builds.""" + return sum(1 for batch in self.embedded_batches if utterance in batch) + class TestSemanticKeywordTierRules: """Test embedding-based keyword_tier_rules matching.""" @@ -1636,6 +1647,42 @@ class TestSemanticKeywordTierRules: # ...and none of it ran on the event-loop thread. assert all(tid != loop_thread_id for tid in fake_router.sync_embedding_thread_ids) + @pytest.mark.asyncio + async def test_concurrent_cold_start_builds_routelayer_once(self, basic_config): + """Concurrent first requests must not each construct the route index (which would + fire duplicate embedding calls); the lazy build happens exactly once. + """ + config = { + **basic_config, + "keyword_tier_rules": [{"keywords": ["kubernetes deployment"], "tier": "REASONING"}], + "semantic_keyword_matching": True, + "embedding_model": "fake-embed", + "match_threshold": 0.5, + } + + def _make_router(fake): + return ComplexityRouter( + model_name="test-router", + litellm_router_instance=fake, + complexity_router_config=config, + ) + + # Baseline: a single cold request's route-index build embeds the route utterance once. + route_utterance = "kubernetes deployment" + baseline_fake = FakeEmbeddingRouter() + await _make_router(baseline_fake)._semantic_tier_override("roll out my k8s cluster", {}) + baseline_builds = baseline_fake.utterance_embedding_count(route_utterance) + assert baseline_builds >= 1 + + # Ten simultaneous cold-start requests must build the index the same number of + # times as one request - i.e. exactly once, not once per concurrent caller. + concurrent_fake = FakeEmbeddingRouter() + concurrent_router = _make_router(concurrent_fake) + await asyncio.gather( + *(concurrent_router._semantic_tier_override("roll out my k8s cluster", {}) for _ in range(10)) + ) + assert concurrent_fake.utterance_embedding_count(route_utterance) == baseline_builds + @pytest.mark.asyncio async def test_below_threshold_falls_back_to_scoring(self, basic_config): """When no route clears the threshold, scoring decides the tier."""