From 72da45d9510adffb40a8052231ebd8545a0f66ad Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 15:38:03 -0700 Subject: [PATCH] address review: clear a failed build via a done-callback, not the waiter The previous cleanup only ran inside a caller's own except handler, so a build that failed after its only caller had already been cancelled left the failed task cached with nothing left to clear it. Move the cleanup onto the task itself as a done-callback, which fires whether or not anyone is still awaiting it, so the next request always gets a fresh attempt instead of replaying the stale failure. Adds a regression test for exactly that ordering (cancel the only caller, let the build fail unobserved, then confirm the next request builds successfully); it fails against the previous except-based cleanup, which left the task cached. --- .../auto_router/auto_router.py | 24 ++++---- .../router_strategy/test_auto_router.py | 59 +++++++++++++++++++ 2 files changed, 72 insertions(+), 11 deletions(-) diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index aaede54e347..7b91fb6db1a 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -138,6 +138,17 @@ class AutoRouter(CustomLogger): self.routelayer = routelayer return routelayer + def _clear_build_task_on_failure(self, build_task: "asyncio.Task[SemanticRouter]") -> None: + """Done-callback: drop a failed build so the next caller gets a fresh attempt. + + Runs whether or not any caller is still awaiting `build_task` (that's the point: + a caller cancelled via `cancel_on_disconnect` mid-build must not leave a later + failure cached with nothing left to clear it), and `not build_task.cancelled()` + guards `.exception()`, which raises on a cancelled task instead of returning one. + """ + if build_task is self._routelayer_build_task and not build_task.cancelled() and build_task.exception(): + self._routelayer_build_task = None + async def _ensure_routelayer(self) -> "SemanticRouter": """Return the cached route layer, building it once under a lock if needed. @@ -158,18 +169,9 @@ class AutoRouter(CustomLogger): build_task = self._routelayer_build_task if build_task is None: build_task = asyncio.ensure_future(asyncio.to_thread(self._build_routelayer)) + build_task.add_done_callback(self._clear_build_task_on_failure) self._routelayer_build_task = build_task - try: - return await asyncio.shield(build_task) - except Exception: - # Only a real build failure (not this caller's own cancellation, which - # asyncio.shield turns into a CancelledError here while the task keeps - # running for everyone else) clears the slot, so the next call retries - # a fresh build instead of replaying the same failure forever. - async with self._routelayer_lock: - if self._routelayer_build_task is build_task: - self._routelayer_build_task = None - raise + return await asyncio.shield(build_task) @staticmethod def _extract_text_from_messages(messages: list[dict[str, Any]]) -> str: diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/test_litellm/router_strategy/test_auto_router.py index cfdff330a23..e8ff99fa7e5 100644 --- a/tests/test_litellm/router_strategy/test_auto_router.py +++ b/tests/test_litellm/router_strategy/test_auto_router.py @@ -721,3 +721,62 @@ class TestAutoRouterColdStartDoesNotBlockTheEventLoop: assert result is not None assert len(embedding_router.embedding_call_threads) == _EMBEDDING_CALLS_PER_ROUTELAYER_BUILD + + @pytest.mark.asyncio + async def test_should_clear_a_failed_build_even_with_no_caller_left_to_observe_it(self): + """Regression: a build that fails after its only caller was already cancelled must + still clear the slot, so the next request gets a fresh attempt instead of replaying + the same stale failure forever.""" + import threading + + class FailsOnFirstAttemptEmbeddingRouter(ThreadTrackingEmbeddingRouter): + def __init__(self) -> None: + super().__init__() + self.started = threading.Event() + self.release = threading.Event() + self.attempts = 0 + + def embedding(self, input: list[str], model: str, **kwargs: Any) -> Any: + self.attempts += 1 + attempt = self.attempts + self.started.set() + self.release.wait(timeout=5) + if attempt == 1: + raise ValueError("boom") + return super().embedding(input, model, **kwargs) + + embedding_router: Final = FailsOnFirstAttemptEmbeddingRouter() + auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router) + + first_call: Final = asyncio.ensure_future( + auto_router.async_pre_routing_hook( + model="my-auto-router", + request_kwargs={}, + messages=[{"role": "user", "content": "fix this stack trace"}], + ) + ) + while not embedding_router.started.is_set(): + await asyncio.sleep(0.01) + first_call.cancel() + with pytest.raises(asyncio.CancelledError): + await first_call + + # Nobody awaits the build now. Let the first attempt fail on its own. + embedding_router.release.set() + build_task = auto_router._routelayer_build_task + assert build_task is not None + while not build_task.done(): + await asyncio.sleep(0.01) + await asyncio.sleep(0.01) # let the done-callback (scheduled via call_soon) run + + assert auto_router._routelayer_build_task is None + + embedding_router.started.clear() + embedding_router.release.clear() + 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