From ba2a704a924ac3236ef2e92d819b882df20f692d Mon Sep 17 00:00:00 2001 From: Vansh Date: Wed, 28 Oct 2026 10:13:45 +0530 Subject: [PATCH] fix(router): use sync Redis pipeline in Router.completion to prevent connection leaks (#48589) - Add synchronous increment_pipeline to RedisCache, DualCache, and InMemoryCache. - Pass sync_mode from Router.function_with_fallbacks down to _increment_deployment_usage. - Use synchronous cache pipeline and local cache set when sync_mode is True, preventing per-call async Redis client creation and event loop socket leaks. - Add regression test asserting that sync Router.completion does not leak async Redis clients. --- litellm/caching/dual_cache.py | 28 ++++++ litellm/caching/in_memory_cache.py | 13 ++- litellm/caching/redis_cache.py | 86 +++++++++++++++++- litellm/integrations/otel/model/spans.py | 1 + litellm/router.py | 51 ++++++++--- .../test_sync_router_redis_connection_leak.py | 87 +++++++++++++++++++ 6 files changed, 253 insertions(+), 13 deletions(-) create mode 100644 tests/unit/test_router/test_sync_router_redis_connection_leak.py diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index e6214a3650c..74dd167e04e 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -615,6 +615,34 @@ class DualCache(BaseCache): _log_deferred_increment_failure ) + def increment_cache_pipeline( + self, + increment_list: Sequence["RedisPipelineIncrementOperation"], + local_only: bool = False, + parent_otel_span: Span | None = None, + **kwargs: object, # kwargs-ok: forward extra params + ) -> list[float] | None: # mutable-ok: list return type + result: list[float] | None = None # rebind-ok: conditional cache assignment # mutable-ok: list result type + try: + result = self.in_memory_cache.increment_pipeline( # rebind-ok: update result from in-memory + increment_list=increment_list, + parent_otel_span=parent_otel_span, + ) + + if self.redis_cache is not None and local_only is False: + result = self.redis_cache.increment_pipeline( # rebind-ok: update result from redis + increment_list=increment_list, + parent_otel_span=parent_otel_span, + ) + except Exception as e: # noqa: BLE001 # fallback to in-memory on redis error + log_redis_failure( + verbose_logger, + logging.WARNING, + "Redis increment_cache_pipeline failed, falling back to in-memory result", + e, + ) + return result + async def async_increment_cache_pipeline( self, increment_list: list["RedisPipelineIncrementOperation"], diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 1cc3e106aa0..4753e053cb8 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -13,7 +13,7 @@ import json import sys import threading import time -from collections.abc import Callable +from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: @@ -269,6 +269,17 @@ class InMemoryCache(BaseCache): async def async_increment(self, key, value: float, **kwargs) -> float: return self.increment_cache(key=key, value=value, **kwargs) + def increment_pipeline( + self, + increment_list: Sequence["RedisPipelineIncrementOperation"], + **kwargs: object, # kwargs-ok: forward extra params + ) -> list[float] | None: # mutable-ok: list return type + results: Final[list[float]] = [ # mutable-ok: list accumulator + self.increment_cache(increment["key"], increment["increment_value"], **kwargs) # pyright: ignore[reportUnknownMemberType] # unannotated base method + for increment in increment_list + ] + return results + async def async_increment_pipeline( self, increment_list: list["RedisPipelineIncrementOperation"], **kwargs ) -> list[float] | None: diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index afe6f8a1eab..11757a06904 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -781,7 +781,6 @@ class RedisCache(BaseCache): def init_async_client( self, ) -> async_redis_client | async_redis_cluster_client: - from redis.asyncio import RedisCluster from litellm import in_memory_llm_clients_cache @@ -792,6 +791,8 @@ class RedisCache(BaseCache): if cached_client is not None: redis_async_client = cast(async_redis_client | async_redis_cluster_client, cached_client) else: + from redis.asyncio import RedisCluster + # Create new connection pool and client for current event loop self.async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs) redis_async_client = get_redis_async_client(connection_pool=self.async_redis_conn_pool, **self.redis_kwargs) @@ -2045,6 +2046,89 @@ class RedisCache(BaseCache): verbose_logger.debug("Increment ASYNC Redis Cache PIPELINE: results: %s", results) return [r for r in results if isinstance(r, float)] + @_redis_circuit_breaker_guard_sync + def increment_pipeline( + self, + increment_list: Sequence[RedisPipelineIncrementOperation], + **kwargs: object, # kwargs-ok: forward extra params + ) -> list[float] | None: # mutable-ok: list return type + """ + Use sync Redis Pipelines for bulk increment operations + Args: + increment_list: List of RedisPipelineIncrementOperation dicts containing: + - key: str + - increment_value: float + - ttl: int | None + """ + if len(increment_list) == 0: + return None + + start_time: Final = time.time() + print_verbose("Increment Sync Redis Cache Pipeline: increment list: %s", increment_list) + + _sync_client: Final = self.redis_client # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # sync client access + if _sync_client is None: # pyright: ignore[reportUnknownMemberType] # sync client check + return None + + try: + with _sync_client.pipeline(transaction=False) as pipe: # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # sync pipeline + incr_indices: Final[list[int]] = [] # mutable-ok: index collection + idx: int = 0 # rebind-ok: pipeline step counter + for increment_op in increment_list: + cache_key = self.check_and_fix_namespace(key=increment_op["key"]) + pipe.incrbyfloat(cache_key, increment_op["increment_value"]) # pyright: ignore[reportUnknownMemberType] # pipeline method + incr_indices.append(idx) + idx += 1 + ttl = increment_op.get("ttl") + if ttl is not None: + _td = timedelta(seconds=float(ttl)) + pipe.expire(cache_key, _td) # pyright: ignore[reportUnknownMemberType] # pipeline method + idx += 1 + results: Final = pipe.execute() # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # pipeline execute + + verbose_logger.debug("Increment SYNC Redis Cache PIPELINE: results: %s", results) # pyright: ignore[reportUnknownArgumentType] # logger call + parsed_results: Final[list[float]] = ( # mutable-ok: list parsed results + [float(results[i]) for i in incr_indices if not isinstance(results[i], bool)] # pyright: ignore[reportUnknownArgumentType] # float cast + if len(results) == idx # pyright: ignore[reportUnknownArgumentType] # sequence length + else [float(r) for r in results if isinstance(r, (float, int)) and not isinstance(r, bool)] # pyright: ignore[reportUnknownArgumentType,reportUnknownVariableType] # float cast + ) + + ## LOGGING ## + end_time: Final = time.time() + _duration: Final = end_time - start_time + self.service_logger_obj.service_success_hook( + service=ServiceTypes.REDIS, + duration=_duration, + call_type="increment_pipeline", + caller=_get_call_stack_info(), + start_time=start_time, + end_time=end_time, + parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), + ) + except Exception as e: + ## LOGGING ## + end_time_err: Final = time.time() + _duration_err: Final = end_time_err - start_time + self.service_logger_obj.service_failure_hook( + service=ServiceTypes.REDIS, + duration=_duration_err, + error=e, + call_type="increment_pipeline", + caller=_get_call_stack_info(), + start_time=start_time, + end_time=end_time_err, + parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), + ) + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: sync increment_pipeline() - Got exception from REDIS", + e, + ) + raise + else: + return parsed_results + @_redis_circuit_breaker_guard async def async_increment_pipeline( self, increment_list: list[RedisPipelineIncrementOperation], **kwargs diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index d8be607150c..f9a6abfe501 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -213,6 +213,7 @@ _SERVICE_VERB_BY_CALL_TYPE: Final[dict[str, str]] = { "async_set_cache_sadd": "sadd", "increment_cache": "incr", "async_increment": "incr", + "increment_pipeline": "incr", "async_increment_pipeline": "incr", "async_increment_if_present": "incr", "delete_cache": "delete", diff --git a/litellm/router.py b/litellm/router.py index c8153b596d6..08d1db7e7c2 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9160,9 +9160,16 @@ class Router: """ model_group: Final = kwargs.get("model") response = original_function(*args, **kwargs) - if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): + is_async_call: Final[bool] = coroutine_checker.is_async_callable(response) or bool( + inspect.isawaitable(response) # pyright: ignore[reportAny] # dynamic response object + ) + if is_async_call: response = await response - await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) + await self.increment_deployment_usage_for_response( + response=response, # pyright: ignore[reportAny] # dynamic response object + request_kwargs=kwargs, # pyright: ignore[reportUnknownArgumentType] # kwargs dict + sync_mode=not is_async_call, + ) ## PROCESS RESPONSE HEADERS response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs) @@ -9379,6 +9386,9 @@ class Router: total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0) counted_tokens: Final = get_counted_usage_tokens(litellm_params) deployment_name: Final = metadata.get("deployment", None) + sync_mode: Final[bool] = isinstance(litellm_params, Mapping) and LiteLLMLogging.is_sync_litellm_request( + litellm_params # pyright: ignore[reportUnknownArgumentType] # mapping from kwargs + ) return await self._increment_deployment_usage( deployment_id=str(deployment_id), deployment_name=deployment_name if isinstance(deployment_name, str) else None, @@ -9386,6 +9396,7 @@ class Router: total_tokens=total_tokens if counted_tokens is None else max(0, total_tokens - counted_tokens), rpm_increment=1 if counted_tokens is None else 0, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), + sync_mode=sync_mode, ) except Exception as e: @@ -9427,6 +9438,7 @@ class Router: self, response: object, request_kwargs: dict[str, object], + sync_mode: bool = False, ) -> None: if response is None: return @@ -9442,6 +9454,10 @@ class Router: total_tokens: Final = response_total_token_count(response) deployment_name: Final = deployment_metadata.get("deployment") deployment_metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] = total_tokens + req_litellm_params: Final = request_kwargs.get("litellm_params") + is_sync: Final[bool] = sync_mode or ( + isinstance(req_litellm_params, Mapping) and LiteLLMLogging.is_sync_litellm_request(req_litellm_params) # pyright: ignore[reportUnknownArgumentType] # mapping from request_kwargs + ) try: await self._increment_deployment_usage( deployment_id=str(deployment_id), @@ -9450,6 +9466,7 @@ class Router: total_tokens=total_tokens, rpm_increment=1, parent_otel_span=get_parent_otel_span_from_kwargs(request_kwargs), + sync_mode=is_sync, ) except Exception: deployment_metadata.pop(ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, None) @@ -9468,6 +9485,7 @@ class Router: total_tokens: float, rpm_increment: int, parent_otel_span: Span | None, + sync_mode: bool = False, ) -> str | None: from litellm.types.caching import RedisPipelineIncrementOperation @@ -9487,15 +9505,26 @@ class Router: RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value) for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment)) ] - post_increment_values: Final = await self.cache.async_increment_cache_pipeline( - increment_list=pipeline_operations, - parent_otel_span=parent_otel_span, - ) - if post_increment_values is not None and self.cache.redis_cache is not None: - for operation, value in zip(pipeline_operations, post_increment_values): - await self.cache.async_set_cache( - operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value - ) + if sync_mode: + sync_post_increment_values: Final = self.cache.increment_cache_pipeline( + increment_list=pipeline_operations, + parent_otel_span=parent_otel_span, + ) + if sync_post_increment_values is not None and self.cache.redis_cache is not None: + for operation, value in zip(pipeline_operations, sync_post_increment_values): + self.cache.set_cache( # pyright: ignore[reportUnknownMemberType] # cache method + operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value + ) + else: + async_post_increment_values: Final = await self.cache.async_increment_cache_pipeline( # pyright: ignore[reportUnknownMemberType] # cache method + increment_list=pipeline_operations, + parent_otel_span=parent_otel_span, + ) + if async_post_increment_values is not None and self.cache.redis_cache is not None: + for operation, value in zip(pipeline_operations, async_post_increment_values): + await self.cache.async_set_cache( # pyright: ignore[reportUnknownMemberType] # cache method + operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value + ) return tpm_key def _deployment_tracks_usage(self, deployment_id: str, model_group: str) -> bool: diff --git a/tests/unit/test_router/test_sync_router_redis_connection_leak.py b/tests/unit/test_router/test_sync_router_redis_connection_leak.py new file mode 100644 index 00000000000..1dcbd655b93 --- /dev/null +++ b/tests/unit/test_router/test_sync_router_redis_connection_leak.py @@ -0,0 +1,87 @@ +""" +Tests to prevent regression for issue #48589: +Sync Router.completion opens a new Redis connection on every call and never closes it +when a deployment has an rpm/tpm limit. +""" + +from unittest.mock import MagicMock, patch + +import litellm +from litellm import Router + + +def test_sync_router_completion_does_not_leak_async_redis_clients() -> None: + """ + Synchronous router.completion() with rpm/tpm limits must use the synchronous + Redis client/pipeline and must NOT instantiate or leak new async Redis clients + into in_memory_llm_clients_cache on each call. + """ + mock_sync_client = MagicMock() + mock_pipeline = MagicMock() + mock_sync_client.pipeline.return_value.__enter__.return_value = mock_pipeline + mock_pipeline.execute.return_value = [1.0, 1.0] + + with ( + patch("litellm._redis.get_redis_client", return_value=mock_sync_client), + patch("litellm._redis.get_redis_async_client", return_value=MagicMock()), + patch("litellm._redis.get_redis_connection_pool", return_value=MagicMock()), + ): + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + "mock_response": "ok", + "rpm": 1000, + }, + } + ], + redis_host="127.0.0.1", + redis_port=6379, + ) + + cache = litellm.in_memory_llm_clients_cache.cache_dict + before_keys = set(cache.keys()) + + for _ in range(5): + response = router.completion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + assert response is not None + + new_async_keys = [k for k in (set(cache.keys()) - before_keys) if k.startswith("async-redis-client")] + assert len(new_async_keys) == 0, f"Leaked async Redis clients in cache: {new_async_keys}" + assert mock_pipeline.incrbyfloat.called, "Sync pipeline incrbyfloat should be called" + assert mock_pipeline.execute.called, "Sync pipeline execute should be called" + + +def test_sync_redis_cache_increment_pipeline_with_real_redis_return_types() -> None: + """ + Real Redis pipeline execution returns [incrbyfloat_val (float), expire_val (bool/int)]. + increment_pipeline must properly extract only the increment results and ignore expire return codes. + """ + from litellm.caching.redis_cache import RedisCache + from litellm.types.caching import RedisPipelineIncrementOperation + + mock_client = MagicMock() + mock_pipeline = MagicMock() + mock_client.pipeline.return_value.__enter__.return_value = mock_pipeline + # Real Redis execute returns [15.5 (float), True (bool), 1.0 (float), True (bool)] + mock_pipeline.execute.return_value = [15.5, True, 1.0, True] + + with patch("litellm._redis.get_redis_client", return_value=mock_client): + cache = RedisCache(host="127.0.0.1", port=6379) + cache.redis_client = mock_client + + ops: list[RedisPipelineIncrementOperation] = [ + {"key": "tpm_key", "increment_value": 15.5, "ttl": 60}, + {"key": "rpm_key", "increment_value": 1.0, "ttl": 60}, + ] + results = cache.increment_pipeline(increment_list=ops) + + assert results == [15.5, 1.0], f"Expected [15.5, 1.0], got {results}" + assert mock_pipeline.incrbyfloat.call_count == 2 + assert mock_pipeline.expire.call_count == 2