fix: Pin redis-py to 4.5.5 and enhance logging for Redis reconnection and failover (#1620)

Author: achimnol

changes/1620.fix.md

@@ -0,0 +1 @@
+Improve logging when retrying redis connections during failover and use explicit names for all redis connection pools

python.lock

@@ -29,7 +29,6 @@ dataclasses-json~=0.5.7
 etcetra==0.1.17
 faker~=13.12.0
 graphene~=2.1.9
-hiredis~=2.2.3
 humanize>=3.1.0
 ifaddr~=0.2
 inquirer~=2.9.2
@@ -55,7 +54,8 @@ pyzmq~=24.0.1
 PyJWT~=2.0
 PyYAML~=6.0
 packaging>=21.3
-redis[hiredis]~=4.6.0
+hiredis>=2.2.3
+redis[hiredis]==4.5.5
 rich~=12.2
 SQLAlchemy[postgresql_asyncpg]~=1.4.40
 setproctitle~=1.3.2
@@ -70,7 +70,7 @@ trafaret~=2.1
 treelib==1.6.1
 typeguard~=2.10
 typing_extensions~=4.3
-uvloop>=0.17; sys_platform != "Windows"
+uvloop~=0.17.0; sys_platform != "Windows"  # 0.18 breaks the API and adds Python 3.12 support
 yarl~=1.8.2  # FIXME: revert to >=1.7 after aio-libs/yarl#862 is resolved
 zipstream-new~=1.1.8
 

requirements.txt

@@ -67,7 +67,7 @@
 
 from ai.backend.common import msgpack, redis_helper
 from ai.backend.common.config import model_definition_iv
-from ai.backend.common.defs import REDIS_STREAM_DB
+from ai.backend.common.defs import REDIS_STAT_DB, REDIS_STREAM_DB
 from ai.backend.common.docker import MAX_KERNELSPEC, MIN_KERNELSPEC, ImageRef
 from ai.backend.common.events import (
     AbstractEvent,
@@ -631,8 +631,16 @@ async def __ainit__(self) -> None:
             node_id=self.local_config["agent"]["id"],
             consumer_group=EVENT_DISPATCHER_CONSUMER_GROUP,
         )
-        self.redis_stream_pool = redis_helper.get_redis_object(self.local_config["redis"], db=4)
-        self.redis_stat_pool = redis_helper.get_redis_object(self.local_config["redis"], db=0)
+        self.redis_stream_pool = redis_helper.get_redis_object(
+            self.local_config["redis"],
+            name="stream",
+            db=REDIS_STREAM_DB,
+        )
+        self.redis_stat_pool = redis_helper.get_redis_object(
+            self.local_config["redis"],
+            name="stat",
+            db=REDIS_STAT_DB,
+        )
 
         alloc_map_mod.log_alloc_map = self.local_config["debug"]["log-alloc-map"]
         computers = await self.load_resources()

src/ai/backend/agent/agent.py

@@ -9,10 +9,8 @@
 import socket
 import uuid
 from collections import defaultdict
-from types import TracebackType
 from typing import (
     Any,
-    Awaitable,
     Callable,
     ClassVar,
     Coroutine,
@@ -32,8 +30,8 @@
 from aiotools.context import aclosing
 from aiotools.server import process_index
 from aiotools.taskgroup import PersistentTaskGroup
+from aiotools.taskgroup.types import AsyncExceptionHandler
 from redis.asyncio import ConnectionPool
-from typing_extensions import TypeAlias
 
 from . import msgpack, redis_helper
 from .logging import BraceStyleAdapter
@@ -60,10 +58,6 @@
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))  # type: ignore[name-defined]
 
-PTGExceptionHandler: TypeAlias = Callable[
-    [Type[Exception], Exception, TracebackType], Awaitable[None]
-]
-
 
 class AbstractEvent(metaclass=abc.ABCMeta):
     # derivatives should define the fields.
@@ -850,16 +844,18 @@ def __init__(
         log_events: bool = False,
         *,
         consumer_group: str,
-        service_name: str = None,
+        service_name: str | None = None,
         stream_key: str = "events",
-        node_id: str = None,
-        consumer_exception_handler: PTGExceptionHandler = None,
-        subscriber_exception_handler: PTGExceptionHandler = None,
+        node_id: str | None = None,
+        consumer_exception_handler: AsyncExceptionHandler | None = None,
+        subscriber_exception_handler: AsyncExceptionHandler | None = None,
     ) -> None:
         _redis_config = redis_config.copy()
         if service_name:
             _redis_config["service_name"] = service_name
-        self.redis_client = redis_helper.get_redis_object(_redis_config, db=db)
+        self.redis_client = redis_helper.get_redis_object(
+            _redis_config, name="event_dispatcher.stream", db=db
+        )
         self._log_events = log_events
         self._closed = False
         self.consumers = defaultdict(set)
@@ -905,7 +901,7 @@ def consume(
         callback: EventCallback[TContext, TEvent],
         coalescing_opts: CoalescingOptions = None,
         *,
-        name: str = None,
+        name: str | None = None,
     ) -> EventHandler[TContext, TEvent]:
         if name is None:
             name = f"evh-{secrets.token_urlsafe(16)}"
@@ -928,9 +924,9 @@ def subscribe(
         event_cls: Type[TEvent],
         context: TContext,
         callback: EventCallback[TContext, TEvent],
-        coalescing_opts: CoalescingOptions = None,
+        coalescing_opts: CoalescingOptions | None = None,
         *,
-        name: str = None,
+        name: str | None = None,
     ) -> EventHandler[TContext, TEvent]:
         if name is None:
             name = f"evh-{secrets.token_urlsafe(16)}"
@@ -1054,15 +1050,19 @@ def __init__(
         redis_config: EtcdRedisConfig,
         db: int = 0,
         *,
-        service_name: str = None,
+        service_name: str | None = None,
         stream_key: str = "events",
         log_events: bool = False,
     ) -> None:
         _redis_config = redis_config.copy()
         if service_name:
             _redis_config["service_name"] = service_name
         self._closed = False
-        self.redis_client = redis_helper.get_redis_object(_redis_config, db=db)
+        self.redis_client = redis_helper.get_redis_object(
+            _redis_config,
+            name="event_producer.stream",
+            db=db,
+        )
         self._log_events = log_events
         self._stream_key = stream_key
 
@@ -1092,7 +1092,7 @@ async def produce_event(
         )
 
 
-def _generate_consumer_id(node_id: str = None) -> str:
+def _generate_consumer_id(node_id: str | None = None) -> str:
     h = hashlib.sha1()
     h.update(str(node_id or socket.getfqdn()).encode("utf8"))
     hostname_hash = h.hexdigest()

src/ai/backend/common/events.py

@@ -4,6 +4,7 @@
 import inspect
 import logging
 import socket
+import time
 from typing import (
     Any,
     AsyncIterator,
@@ -21,7 +22,7 @@
 
 import redis.exceptions
 import yarl
-from redis.asyncio import Redis
+from redis.asyncio import ConnectionPool, Redis
 from redis.asyncio.client import Pipeline, PubSub
 from redis.asyncio.sentinel import MasterNotFoundError, Sentinel, SlaveNotFoundError
 from redis.backoff import ExponentialBackoff
@@ -63,6 +64,10 @@
         redis.exceptions.TimeoutError,
     ],
 }
+_default_conn_pool_opts: Mapping[str, Any] = {
+    "max_connections": 16,
+    # "timeout": 20.0,  # for redis-py 5.0+
+}
 
 _scripts: Dict[str, str] = {}
 
@@ -73,11 +78,6 @@ class ConnectionNotAvailable(Exception):
     pass
 
 
-def _parse_stream_msg_id(msg_id: bytes) -> Tuple[int, int]:
-    timestamp, _, sequence = msg_id.partition(b"-")
-    return int(timestamp), int(sequence)
-
-
 async def subscribe(channel: PubSub, *, reconnect_poll_interval: float = 0.3) -> AsyncIterator[Any]:
     """
     An async-generator wrapper for pub-sub channel subscription.
@@ -130,7 +130,7 @@ async def blpop(
     redis_obj: RedisConnectionInfo,
     key: str,
     *,
-    service_name: str = None,
+    service_name: Optional[str] = None,
 ) -> AsyncIterator[Any]:
     """
     An async-generator wrapper for blpop (blocking left pop).
@@ -175,8 +175,9 @@ async def execute(
     redis_obj: RedisConnectionInfo,
     func: Callable[[Redis], Awaitable[Any]],
     *,
-    service_name: str = None,
+    service_name: Optional[str] = None,
     encoding: Optional[str] = None,
+    command_timeout: Optional[float] = None,
 ) -> Any:
     """
     Executes a function that issues Redis commands or returns a pipeline/transaction of commands,
@@ -187,9 +188,25 @@ async def execute(
     """
     redis_client = redis_obj.client
     service_name = service_name or redis_obj.service_name
-    reconnect_poll_interval = float(
-        cast(str, redis_obj.redis_helper_config.get("reconnect_poll_timeout"))
-    )
+    reconnect_poll_interval = redis_obj.redis_helper_config.get("reconnect_poll_timeout", 0.0)
+
+    first_trial = time.perf_counter()
+    retry_log_count = 0
+    last_log_time = first_trial
+
+    def show_retry_warning(e: Exception, warn_on_first_attempt: bool = True) -> None:
+        nonlocal retry_log_count, last_log_time
+        now = time.perf_counter()
+        if (warn_on_first_attempt and retry_log_count == 0) or now - last_log_time >= 10.0:
+            log.warning(
+                "Retrying due to interruption of Redis connection "
+                "({}, conn-pool: {}, retrying-for: {:.3f}s)",
+                repr(e),
+                redis_obj.name,
+                now - first_trial,
+            )
+            retry_log_count += 1
+            last_log_time = now
 
     while True:
         try:
@@ -228,21 +245,33 @@ async def execute(
             MasterNotFoundError,
             SlaveNotFoundError,
             redis.exceptions.ReadOnlyError,
+            redis.exceptions.ConnectionError,
             ConnectionResetError,
-        ):
+        ) as e:
+            warn_on_first_attempt = True
+            if (
+                isinstance(e, redis.exceptions.ConnectionError)
+                and "Too many connections" in e.args[0]
+            ):  # connection pool is full
+                warn_on_first_attempt = False
+            show_retry_warning(e, warn_on_first_attempt)
             await asyncio.sleep(reconnect_poll_interval)
             continue
-        except redis.exceptions.ConnectionError as e:
-            log.error(f"execute(): Connecting to redis failed: {e}")
-            await asyncio.sleep(reconnect_poll_interval)
+        except (
+            redis.exceptions.TimeoutError,
+            asyncio.TimeoutError,
+        ) as e:
+            if command_timeout is not None:
+                now = time.perf_counter()
+                if now - first_trial >= command_timeout + 1.0:
+                    show_retry_warning(e)
             continue
         except redis.exceptions.ResponseError as e:
             if "NOREPLICAS" in e.args[0]:
+                show_retry_warning(e)
                 await asyncio.sleep(reconnect_poll_interval)
                 continue
             raise
-        except (redis.exceptions.TimeoutError, asyncio.TimeoutError):
-            continue
         except asyncio.CancelledError:
             raise
         finally:
@@ -317,6 +346,7 @@ async def read_stream(
                     {stream_key: last_id},
                     block=block_timeout,
                 ),
+                command_timeout=block_timeout / 1000,
             )
             if not reply:
                 continue
@@ -367,6 +397,7 @@ async def read_stream_by_group(
                         str(autoclaim_idle_timeout),
                         autoclaim_start_id,
                     ),
+                    command_timeout=autoclaim_idle_timeout / 1000,
                 )
                 for msg_id, msg_data in reply[1]:
                     messages.append((msg_id, msg_data))
@@ -381,6 +412,7 @@ async def read_stream_by_group(
                     {stream_key: b">"},  # fetch messages not seen by other consumers
                     block=block_timeout,
                 ),
+                command_timeout=block_timeout / 1000,
             )
             if len(reply) == 0:
                 continue
@@ -422,13 +454,32 @@ async def read_stream_by_group(
 
 def get_redis_object(
     redis_config: EtcdRedisConfig,
+    *,
+    name: str,
     db: int = 0,
     **kwargs,
 ) -> RedisConnectionInfo:
     redis_helper_config: RedisHelperConfig = cast(
         RedisHelperConfig, redis_config.get("redis_helper_config")
     )
-
+    conn_opts = {
+        **_default_conn_opts,
+        **kwargs,
+        # "lib_name": None,  # disable implicit "CLIENT SETINFO" (for redis-py 5.0+)
+        # "lib_version": None,  # disable implicit "CLIENT SETINFO" (for redis-py 5.0+)
+    }
+    conn_pool_opts = {
+        **_default_conn_pool_opts,
+    }
+    if socket_timeout := redis_helper_config.get("socket_timeout"):
+        conn_opts["socket_timeout"] = float(socket_timeout)
+    if socket_connect_timeout := redis_helper_config.get("socket_connect_timeout"):
+        conn_opts["socket_connect_timeout"] = float(socket_connect_timeout)
+    if max_connections := redis_helper_config.get("max_connections"):
+        conn_pool_opts["max_connections"] = int(max_connections)
+    # for redis-py 5.0+
+    # if connection_ready_timeout := redis_helper_config.get("connection_ready_timeout"):
+    #     conn_pool_opts["timeout"] = float(connection_ready_timeout)
     if _sentinel_addresses := redis_config.get("sentinel"):
         sentinel_addresses: Any = None
         if isinstance(_sentinel_addresses, str):
@@ -453,19 +504,14 @@ def get_redis_object(
                 **kwargs,
             },
         )
-
-        conn_opts = {
-            **_default_conn_opts,
-            **kwargs,
-            "socket_timeout": float(cast(str, redis_helper_config.get("socket_timeout"))),
-            "socket_connect_timeout": float(
-                cast(str, redis_helper_config.get("socket_connect_timeout"))
-            ),
-        }
-
         return RedisConnectionInfo(
-            client=sentinel.master_for(service_name=service_name, password=password, **conn_opts),
+            client=sentinel.master_for(
+                service_name=service_name,
+                password=password,
+                **conn_opts,
+            ),
             sentinel=sentinel,
+            name=name,
             service_name=service_name,
             redis_helper_config=redis_helper_config,
         )
@@ -475,10 +521,18 @@ def get_redis_object(
         url = yarl.URL("redis://host").with_host(str(redis_url[0])).with_port(
             redis_url[1]
         ).with_password(redis_config.get("password")) / str(db)
-
         return RedisConnectionInfo(
-            client=Redis.from_url(str(url), **kwargs),
+            # In redis-py 5.0.1+, we should migrate to `Redis.from_pool()` API
+            client=Redis(
+                connection_pool=ConnectionPool.from_url(
+                    str(url),
+                    **conn_pool_opts,
+                ),
+                **conn_opts,
+                auto_close_connection_pool=True,
+            ),
             sentinel=None,
+            name=name,
             service_name=None,
             redis_helper_config=redis_helper_config,
         )