fix(router): give each least-busy attempt its own in-flight member when call ids repeat

Co-Authored-By: Krishnan <keshi8086@gmail.com>

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)

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,