From 5a2845d183d588fa892c7409b2fed077fd568c3d Mon Sep 17 00:00:00 2001 From: moe-berri Date: Thu, 3 Sep 2026 20:10:02 -0700 Subject: [PATCH] fix(router): preserve classifier breaker state under concurrency --- .../complexity_router/complexity_router.py | 47 ++++++++---- .../router_strategy/test_complexity_router.py | 75 +++++++++++++++---- 2 files changed, 94 insertions(+), 28 deletions(-) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 9176b3da02a..d7bbf85d4fa 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -836,31 +836,45 @@ class _ClassifierCircuitBreaker: self._clock = clock self._state = self.CLOSED self._opened_at: float | None = None + self._generation = 0 self._lock = Lock() - def allow_request(self) -> bool: - """Allow ordinary calls while closed and exactly one probe after cooldown.""" + def acquire_permit(self) -> int | None: + """Return a generation-scoped permit, or deny the call while the circuit is open. + + Calls admitted together while closed share a generation. The first timeout advances it, + making every other in-flight completion stale so it cannot erase the new cooldown. + """ with self._lock: if self._state == self.CLOSED: - return True + return self._generation if self._state == self.HALF_OPEN: - return False + return None opened_at: Final = self._opened_at if opened_at is not None and self._clock() - opened_at >= self._cooldown_seconds: self._state = self.HALF_OPEN - return True - return False + return self._generation + return None - def record_success(self) -> None: + def record_success(self, permit: int) -> None: + """Close only when the current half-open recovery probe succeeds.""" with self._lock: + if self._state != self.HALF_OPEN or permit != self._generation: + return self._state = self.CLOSED self._opened_at = None - def record_failure(self, *, is_timeout: bool) -> None: + def record_failure(self, permit: int, *, is_timeout: bool) -> None: """Open on a normal timeout, or reopen when the single recovery probe fails.""" with self._lock: - if not is_timeout and self._state != self.HALF_OPEN: + if permit != self._generation: return + if self._state == self.CLOSED: + if not is_timeout: + return + elif self._state != self.HALF_OPEN: + return + self._generation += 1 self._state = self.OPEN self._opened_at = self._clock() @@ -1541,7 +1555,8 @@ class ComplexityRouter(CustomLogger): has. It is handed to the failure path so a classifier error does not re-run the scorer. """ breaker: Final = self._classifier_circuit_breaker - if breaker is not None and not breaker.allow_request(): + permit: Final = breaker.acquire_permit() if breaker is not None else None + if breaker is not None and permit is None: return self._classifier_failure_outcome( "LLM classifier circuit is open", prompt, @@ -1551,8 +1566,8 @@ class ComplexityRouter(CustomLogger): ) try: tier, classifier_cost = await self._classify_with_llm(prompt, system_prompt, request_kwargs, messages) - if breaker is not None: - breaker.record_success() + if breaker is not None and permit is not None: + breaker.record_success(permit) return ClassificationOutcome( tier=tier, score=None, @@ -1560,9 +1575,13 @@ class ComplexityRouter(CustomLogger): cause="llm_classifier", classifier_cost=classifier_cost, ) + except asyncio.CancelledError: + if breaker is not None and permit is not None: + breaker.record_failure(permit, is_timeout=False) + raise except Exception as e: # noqa: BLE001 -- external LLM call can fail in many distinct ways (timeout, provider error, validation, parse error); any failure must fall back to the configured fallback path - if breaker is not None: - breaker.record_failure(is_timeout=_is_classifier_timeout(e)) + if breaker is not None and permit is not None: + breaker.record_failure(permit, is_timeout=_is_classifier_timeout(e)) return self._classifier_failure_outcome(f"LLM classifier failed ({e})", prompt, system_prompt, scored) def _classifier_failure_outcome( diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index d55cabd6806..137fb128a57 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -2100,29 +2100,74 @@ class TestLLMClassifier: now = 100.0 breaker = _ClassifierCircuitBreaker(30.0, clock=lambda: now) - assert breaker.allow_request() is True - breaker.record_failure(is_timeout=True) - assert breaker.allow_request() is False + initial_permit = breaker.acquire_permit() + assert initial_permit is not None + breaker.record_failure(initial_permit, is_timeout=True) + assert breaker.acquire_permit() is None now = 130.0 - assert breaker.allow_request() is True - assert breaker.allow_request() is False + probe_permit = breaker.acquire_permit() + assert probe_permit is not None + assert breaker.acquire_permit() is None - breaker.record_success() - assert breaker.allow_request() is True + breaker.record_success(probe_permit) + assert breaker.acquire_permit() is not None def test_failed_classifier_probe_restarts_cooldown(self): now = 100.0 breaker = _ClassifierCircuitBreaker(30.0, clock=lambda: now) - breaker.record_failure(is_timeout=True) + initial_permit = breaker.acquire_permit() + assert initial_permit is not None + breaker.record_failure(initial_permit, is_timeout=True) now = 130.0 - assert breaker.allow_request() is True - breaker.record_failure(is_timeout=False) - assert breaker.allow_request() is False + probe_permit = breaker.acquire_permit() + assert probe_permit is not None + breaker.record_failure(probe_permit, is_timeout=False) + assert breaker.acquire_permit() is None now = 160.0 - assert breaker.allow_request() is True + assert breaker.acquire_permit() is not None + + def test_stale_success_cannot_close_circuit_opened_by_overlapping_timeout(self): + breaker = _ClassifierCircuitBreaker(30.0) + timeout_permit = breaker.acquire_permit() + stale_success_permit = breaker.acquire_permit() + assert timeout_permit is not None + assert stale_success_permit is not None + + breaker.record_failure(timeout_permit, is_timeout=True) + breaker.record_success(stale_success_permit) + + assert breaker.acquire_permit() is None + + @pytest.mark.asyncio + async def test_cancelled_classifier_probe_restarts_cooldown(self, mock_router_instance, llm_classifier_config): + now = 100.0 + mock_router_instance.acompletion = AsyncMock( + side_effect=[ + TimeoutError("classifier timed out"), + asyncio.CancelledError(), + _llm_response('{"tier": "SIMPLE"}'), + ] + ) + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=llm_classifier_config, + ) + router._classifier_circuit_breaker = _ClassifierCircuitBreaker(30.0, clock=lambda: now) + + await router.aclassify("open the circuit") + now = 130.0 + with pytest.raises(asyncio.CancelledError): + await router.aclassify("cancel the recovery probe") + + outcome = await router.aclassify("stay in cooldown") + + assert outcome.cause == "heuristic_scorer" + assert "classifier-circuit-open" in outcome.signals + assert mock_router_instance.acompletion.await_count == 2 @pytest.mark.asyncio async def test_classifier_circuit_can_be_disabled(self, mock_router_instance, llm_classifier_config): @@ -2146,8 +2191,10 @@ class TestLLMClassifier: def test_non_timeout_failure_does_not_open_closed_classifier_circuit(self): breaker = _ClassifierCircuitBreaker(30.0) - breaker.record_failure(is_timeout=False) - assert breaker.allow_request() is True + permit = breaker.acquire_permit() + assert permit is not None + breaker.record_failure(permit, is_timeout=False) + assert breaker.acquire_permit() is not None @pytest.mark.asyncio async def test_aclassify_classifier_cost_is_none_when_call_is_unpriced(