From 7e25fa409e81fd95dd1dba9c9fd93d72459bf9ab Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 04:03:19 +0000 Subject: [PATCH] fix(router): give each least-busy attempt its own in-flight member when call ids repeat Co-Authored-By: Krishnan --- litellm/router_strategy/least_busy.py | 27 ++++++++++++++----- tests/unit/router_strategy/test_least_busy.py | 24 +++++++++++++++++ 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 75899274eb2..8efbecbb31d 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,5 +1,6 @@ import logging import time +import uuid from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from typing import Final @@ -133,6 +134,20 @@ class LeastBusyLoggingHandler(CustomLogger): self.router_cache = router_cache self.router_cache_id = str(id(router_cache)) self._clock = clock + self._worker_id = uuid.uuid4().hex + + def _attempt_member(self, shared_key: str, call_id: str, delta: int) -> str: + """Name this call's sorted-set member, unique per concurrent attempt sharing a ``litellm_call_id``. + + Callers can supply their own call id, so two concurrent calls may share one. Numbering the + attempts per worker keeps one member per call, and finishing either call removes only one. + """ + attempts_key: Final = f"{shared_key}:{call_id}:attempts" + memory: Final = self.router_cache.in_memory_cache + after: Final = int(memory.increment_cache(attempts_key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)) # pyright: ignore[reportUnknownMemberType] # InMemoryCache.increment_cache leaves key untyped + if after <= 0: + memory.delete_cache(attempts_key) + return f"{call_id}:{self._worker_id}:{after if delta > 0 else after + 1}" @with_service_target("router_usage") def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: @@ -225,10 +240,11 @@ class LeastBusyLoggingHandler(CustomLogger): if redis_cache is None or ref.call_id is None: return shared_key: Final = _in_flight_key(ref.model_group, ref.deployment_id) + member: Final = self._attempt_member(shared_key, ref.call_id, delta) if delta > 0: - redis_cache.track_in_flight(shared_key, ref.call_id, self._clock(), IN_FLIGHT_COUNT_TTL_SECONDS) + redis_cache.track_in_flight(shared_key, member, self._clock(), IN_FLIGHT_COUNT_TTL_SECONDS) else: - redis_cache.untrack_in_flight(shared_key, ref.call_id) + redis_cache.untrack_in_flight(shared_key, member) except Exception as e: _warn_unwritable(key, e) @@ -247,11 +263,10 @@ class LeastBusyLoggingHandler(CustomLogger): if redis_cache is None or ref.call_id is None: return shared_key: Final = _in_flight_key(ref.model_group, ref.deployment_id) + member: Final = self._attempt_member(shared_key, ref.call_id, delta) if delta > 0: - await redis_cache.async_track_in_flight( - shared_key, ref.call_id, self._clock(), IN_FLIGHT_COUNT_TTL_SECONDS - ) + await redis_cache.async_track_in_flight(shared_key, member, self._clock(), IN_FLIGHT_COUNT_TTL_SECONDS) else: - await redis_cache.async_untrack_in_flight(shared_key, ref.call_id) + await redis_cache.async_untrack_in_flight(shared_key, member) except Exception as e: _warn_unwritable(key, e) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 496b7b97380..dc9a7752040 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -145,6 +145,30 @@ async def test_a_request_outliving_the_ttl_cannot_cancel_a_newer_request_when_it assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B +@pytest.mark.asyncio +async def test_finishing_one_of_several_calls_sharing_a_call_id_keeps_the_others_counted( + shared_redis: RedisCache, +) -> None: + """Callers can pass their own litellm_call_id, so concurrent calls may share one. Tracking the + id itself merged them into one entry, and the first to finish hid the rest from other workers""" + worker: Final = _worker(shared_redis) + other_worker: Final = _worker(shared_redis) + picking_worker: Final = _worker(shared_redis) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "shared")) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "shared")) + other_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "shared")) + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "b")) + await worker.async_log_success_event(_call_kwargs("dep-a", "shared"), None, None, None) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + worker.log_success_event(_call_kwargs("dep-a", "shared"), None, None, None) + await other_worker.async_log_failure_event(_call_kwargs("dep-a", "shared"), None, None, None) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + @pytest.mark.asyncio async def test_a_request_nobody_finished_stops_counting_once_it_is_older_than_the_ttl( shared_redis: RedisCache,