fix(model-engine): harden aioredis pool against stale-socket errors (#816)

Author: olliestanley

model-engine/model_engine_server/api/dependencies.py

@@ -7,6 +7,7 @@
 import redis.asyncio as aioredis
 from fastapi import Depends, HTTPException, status
 from fastapi.security import HTTPBasic, HTTPBasicCredentials, OAuth2PasswordBearer
+from model_engine_server.common.aioredis_pool import build_aioredis_pool
 from model_engine_server.common.config import hmi_config
 from model_engine_server.common.dtos.model_endpoints import BrokerType
 from model_engine_server.common.env_vars import CIRCLECI
@@ -560,6 +561,6 @@ def get_or_create_aioredis_pool() -> aioredis.ConnectionPool:
 
     expiration_timestamp = hmi_config.cache_redis_url_expiration_timestamp
     if _pool is None or (expiration_timestamp is not None and time.time() > expiration_timestamp):
-        _pool = aioredis.BlockingConnectionPool.from_url(hmi_config.cache_redis_url)
+        _pool = build_aioredis_pool(hmi_config.cache_redis_url)
     assert _pool is not None
     return _pool

model-engine/model_engine_server/common/aioredis_pool.py

@@ -0,0 +1,52 @@
+"""Shared construction for aioredis connection pools and clients.
+
+redis-py defaults leave pooled connections vulnerable to middlebox idle
+timeouts (Istio sidecars, cloud load balancers, NAT): when `retry` and
+`retry_on_error` are both unset, `AbstractConnection.__init__` uses
+`Retry(NoBackoff(), 0)` — zero retries — so the first command issued on a
+silently-closed pooled socket surfaces as `ConnectionError` to the caller.
+
+The helpers here turn on TCP keepalive, a pre-command PING after idle periods,
+and transparent retry with exponential backoff. Use them for every long-lived
+aioredis pool or client in this repo.
+"""
+
+from typing import Any, Dict
+
+import redis.asyncio as aioredis
+from redis.asyncio.retry import Retry
+from redis.backoff import ExponentialBackoff
+from redis.exceptions import ConnectionError as RedisConnectionError
+from redis.exceptions import TimeoutError as RedisTimeoutError
+
+HEALTH_CHECK_INTERVAL_SECONDS = 30
+SOCKET_CONNECT_TIMEOUT_SECONDS = 5
+RETRY_BACKOFF_CAP_SECONDS = 2.0
+RETRY_BACKOFF_BASE_SECONDS = 0.2
+RETRY_MAX_ATTEMPTS = 3
+
+
+def get_aioredis_connection_kwargs() -> Dict[str, Any]:
+    """Return kwargs for constructing a resilient aioredis pool or client.
+
+    A fresh Retry and error list are returned on every call so callers can
+    never accidentally share mutable state across pools.
+    """
+    return {
+        "socket_keepalive": True,
+        "socket_connect_timeout": SOCKET_CONNECT_TIMEOUT_SECONDS,
+        "health_check_interval": HEALTH_CHECK_INTERVAL_SECONDS,
+        "retry_on_error": [RedisConnectionError, RedisTimeoutError],
+        "retry": Retry(
+            ExponentialBackoff(cap=RETRY_BACKOFF_CAP_SECONDS, base=RETRY_BACKOFF_BASE_SECONDS),
+            retries=RETRY_MAX_ATTEMPTS,
+        ),
+    }
+
+
+def build_aioredis_pool(url: str) -> aioredis.ConnectionPool:
+    return aioredis.BlockingConnectionPool.from_url(url, **get_aioredis_connection_kwargs())
+
+
+def build_aioredis_client(url: str) -> aioredis.Redis:
+    return aioredis.from_url(url, **get_aioredis_connection_kwargs())

model-engine/model_engine_server/core/celery/app.py

@@ -9,6 +9,7 @@
 from celery.app import backends
 from celery.app.control import Inspect
 from celery.result import AsyncResult
+from model_engine_server.common.aioredis_pool import build_aioredis_client
 from model_engine_server.core.aws.roles import session
 from model_engine_server.core.aws.secrets import get_key_file
 from model_engine_server.core.config import infra_config
@@ -232,7 +233,7 @@ def get_redis_instance(db_index: int = 0) -> Union[Redis, StrictRedis]:
 
 def get_async_redis_instance(db_index: int = 0) -> aioredis.Redis:
     host, port = get_redis_host_port()
-    return aioredis.Redis.from_url(f"redis://{host}:{port}/{db_index}")
+    return build_aioredis_client(f"redis://{host}:{port}/{db_index}")
 
 
 def celery_app(

model-engine/model_engine_server/core/celery/celery_autoscaler.py

@@ -11,7 +11,6 @@
 from math import ceil
 from typing import Any, DefaultDict, Dict, List, Set, Tuple
 
-import redis.asyncio as aioredis
 import stringcase
 from azure.core.exceptions import ResourceNotFoundError
 from azure.identity import DefaultAzureCredential
@@ -22,6 +21,7 @@
 from kubernetes_asyncio import config as kube_config
 from kubernetes_asyncio.client.rest import ApiException
 from kubernetes_asyncio.config.config_exception import ConfigException
+from model_engine_server.common.aioredis_pool import build_aioredis_client
 from model_engine_server.core.aws.roles import session
 from model_engine_server.core.celery import (
     TaskVisibility,
@@ -350,7 +350,7 @@ async def _init_client(self):
             get_redis_host_port()
         )  # Switches the redis instance based on CELERY_ELASTICACHE_ENABLED's value
         self.redis = {
-            db_index: aioredis.Redis.from_url(f"redis://{host}:{port}/{db_index}")
+            db_index: build_aioredis_client(f"redis://{host}:{port}/{db_index}")
             for db_index in get_all_db_indexes()
         }
         self.initialized = True

model-engine/model_engine_server/entrypoints/start_batch_job_orchestration.py

@@ -6,6 +6,7 @@
 
 import redis.asyncio as aioredis
 from model_engine_server.api.dependencies import get_monitoring_metrics_gateway
+from model_engine_server.common.aioredis_pool import build_aioredis_pool
 from model_engine_server.common.config import hmi_config
 from model_engine_server.common.dtos.model_endpoints import BrokerType
 from model_engine_server.common.env_vars import CIRCLECI
@@ -65,7 +66,7 @@ async def run_batch_job(
 ):
     tracing_gateway = get_tracing_gateway()
     session = get_session_async_null_pool()
-    pool = aioredis.BlockingConnectionPool.from_url(hmi_config.cache_redis_url)
+    pool = build_aioredis_pool(hmi_config.cache_redis_url)
     redis: aioredis.Redis[Any] = aioredis.Redis(connection_pool=pool)
     sqs_task_queue_gateway = CeleryTaskQueueGateway(
         broker_type=BrokerType.SQS, tracing_gateway=tracing_gateway

model-engine/model_engine_server/infra/gateways/redis_inference_autoscaling_metrics_gateway.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 import redis.asyncio as aioredis
+from model_engine_server.common.aioredis_pool import build_aioredis_client
 from model_engine_server.domain.gateways.inference_autoscaling_metrics_gateway import (
     InferenceAutoscalingMetricsGateway,
 )
@@ -19,7 +20,7 @@ def __init__(
             # If aioredis cannot create a connection pool, reraise that as an error because the
             # default error message is cryptic and not obvious.
             try:
-                self._redis = aioredis.from_url(redis_info, health_check_interval=60)
+                self._redis = build_aioredis_client(redis_info)
             except Exception as exc:
                 raise RuntimeError(
                     "If redis_info is specified, RedisInferenceAutoscalingMetricsGateway must be"

model-engine/model_engine_server/infra/repositories/redis_feature_flag_repository.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 import redis.asyncio as aioredis
+from model_engine_server.common.aioredis_pool import build_aioredis_client
 from model_engine_server.infra.repositories.feature_flag_repository import FeatureFlagRepository
 
 
@@ -15,7 +16,7 @@ def __init__(
             # If aioredis cannot create a connection pool, reraise that as an error because the
             # default error message is cryptic and not obvious.
             try:
-                self._redis = aioredis.from_url(redis_info, health_check_interval=60)
+                self._redis = build_aioredis_client(redis_info)
             except Exception as exc:
                 raise RuntimeError(
                     "If redis_info is specified, RedisFeatureFlagRepository must be"

model-engine/model_engine_server/infra/repositories/redis_model_endpoint_cache_repository.py

@@ -3,6 +3,7 @@
 from typing import Optional
 
 import redis.asyncio as aioredis
+from model_engine_server.common.aioredis_pool import build_aioredis_client
 from model_engine_server.domain.entities import ModelEndpointInfraState
 from model_engine_server.infra.repositories.model_endpoint_cache_repository import (
     ModelEndpointCacheRepository,
@@ -23,7 +24,7 @@ def __init__(
             # If aioredis cannot create a connection pool, reraise that as an error because the
             # default error message is cryptic and not obvious.
             try:
-                self._redis = aioredis.from_url(redis_info, health_check_interval=60)
+                self._redis = build_aioredis_client(redis_info)
             except Exception as exc:
                 raise RuntimeError(
                     "If redis_info is specified, RedisModelEndpointCacheRepository must be"

model-engine/model_engine_server/service_builder/tasks_v1.py

@@ -8,6 +8,7 @@
 from celery.signals import worker_process_init
 from celery.utils.log import get_task_logger
 from model_engine_server.api.dependencies import get_monitoring_metrics_gateway
+from model_engine_server.common.aioredis_pool import build_aioredis_pool
 from model_engine_server.common.config import hmi_config
 from model_engine_server.common.constants import READYZ_FPATH
 from model_engine_server.common.dtos.endpoint_builder import (
@@ -155,7 +156,7 @@ async def _build_endpoint(
         # Redis connection
         if infra_config().debug_mode:  # pragma: no cover
             logger.info("Connecting to Redis", extra={"redis_url": hmi_config.cache_redis_url})
-        pool = aioredis.BlockingConnectionPool.from_url(hmi_config.cache_redis_url)
+        pool = build_aioredis_pool(hmi_config.cache_redis_url)
         redis = aioredis.Redis(connection_pool=pool)
         if infra_config().debug_mode:  # pragma: no cover
             logger.info("Redis connection established successfully")

model-engine/tests/unit/common/test_aioredis_pool.py

@@ -0,0 +1,49 @@
+import redis.asyncio as aioredis
+from model_engine_server.common.aioredis_pool import (
+    HEALTH_CHECK_INTERVAL_SECONDS,
+    RETRY_MAX_ATTEMPTS,
+    SOCKET_CONNECT_TIMEOUT_SECONDS,
+    build_aioredis_client,
+    build_aioredis_pool,
+    get_aioredis_connection_kwargs,
+)
+from redis.asyncio.retry import Retry
+from redis.exceptions import ConnectionError as RedisConnectionError
+from redis.exceptions import TimeoutError as RedisTimeoutError
+
+
+def test_kwargs_include_keepalive_health_check_and_retry():
+    kwargs = get_aioredis_connection_kwargs()
+    assert kwargs["socket_keepalive"] is True
+    assert kwargs["socket_connect_timeout"] == SOCKET_CONNECT_TIMEOUT_SECONDS
+    assert kwargs["health_check_interval"] == HEALTH_CHECK_INTERVAL_SECONDS
+    assert kwargs["retry_on_error"] == [RedisConnectionError, RedisTimeoutError]
+    assert isinstance(kwargs["retry"], Retry)
+
+
+def test_kwargs_do_not_share_mutable_state_across_calls():
+    first = get_aioredis_connection_kwargs()
+    second = get_aioredis_connection_kwargs()
+    assert first["retry"] is not second["retry"]
+    assert first["retry_on_error"] is not second["retry_on_error"]
+
+
+def test_build_aioredis_pool_applies_kwargs_to_connection():
+    pool = build_aioredis_pool("redis://localhost:6379/0")
+    assert isinstance(pool, aioredis.BlockingConnectionPool)
+    conn = pool.connection_class(**pool.connection_kwargs)
+    assert conn.socket_keepalive is True
+    assert conn.socket_connect_timeout == SOCKET_CONNECT_TIMEOUT_SECONDS
+    assert conn.health_check_interval == HEALTH_CHECK_INTERVAL_SECONDS
+    # retries is the only reliable way to tell we replaced the default
+    # NoBackoff(0) with our configured Retry.
+    assert conn.retry._retries == RETRY_MAX_ATTEMPTS
+
+
+def test_build_aioredis_client_applies_kwargs_to_connection():
+    client = build_aioredis_client("redis://localhost:6379/0")
+    pool = client.connection_pool
+    conn = pool.connection_class(**pool.connection_kwargs)
+    assert conn.socket_keepalive is True
+    assert conn.health_check_interval == HEALTH_CHECK_INTERVAL_SECONDS
+    assert conn.retry._retries == RETRY_MAX_ATTEMPTS