From 8ff65abbc6631530f2aaa54584b8f941c13cc2f2 Mon Sep 17 00:00:00 2001 From: liuhaoran <75237518+liuhaoran1212@users.noreply.github.com> Date: Wed, 20 Nov 2024 13:44:35 +0800 Subject: [PATCH] ext_redis.py support redis clusters --- Fixes #9538 (#9789) Signed-off-by: root Co-authored-by: root Co-authored-by: Bowen Liang --- api/.env.example | 5 ++++ api/configs/middleware/cache/redis_config.py | 15 ++++++++++++ api/extensions/ext_redis.py | 9 ++++++- .../unit_tests/core/test_model_manager.py | 24 ++++++++++++------- docker/.env.example | 6 +++++ docker/docker-compose.yaml | 3 +++ 6 files changed, 52 insertions(+), 10 deletions(-) diff --git a/api/.env.example b/api/.env.example index 1a242a3daa..f8a2812563 100644 --- a/api/.env.example +++ b/api/.env.example @@ -42,6 +42,11 @@ REDIS_SENTINEL_USERNAME= REDIS_SENTINEL_PASSWORD= REDIS_SENTINEL_SOCKET_TIMEOUT=0.1 +# redis Cluster configuration. +REDIS_USE_CLUSTERS=false +REDIS_CLUSTERS= +REDIS_CLUSTERS_PASSWORD= + # PostgreSQL database configuration DB_USERNAME=postgres DB_PASSWORD=difyai123456 diff --git a/api/configs/middleware/cache/redis_config.py b/api/configs/middleware/cache/redis_config.py index 26b9b1347c..2e98c31ec3 100644 --- a/api/configs/middleware/cache/redis_config.py +++ b/api/configs/middleware/cache/redis_config.py @@ -68,3 +68,18 @@ class RedisConfig(BaseSettings): description="Socket timeout in seconds for Redis Sentinel connections", default=0.1, ) + + REDIS_USE_CLUSTERS: bool = Field( + description="Enable Redis Clusters mode for high availability", + default=False, + ) + + REDIS_CLUSTERS: Optional[str] = Field( + description="Comma-separated list of Redis Clusters nodes (host:port)", + default=None, + ) + + REDIS_CLUSTERS_PASSWORD: Optional[str] = Field( + description="Password for Redis Clusters authentication (if required)", + default=None, + ) diff --git a/api/extensions/ext_redis.py b/api/extensions/ext_redis.py index e1f8409f21..36f06c1104 100644 --- a/api/extensions/ext_redis.py +++ b/api/extensions/ext_redis.py @@ -1,11 +1,12 @@ import redis +from redis.cluster import ClusterNode, RedisCluster from redis.connection import Connection, SSLConnection from redis.sentinel import Sentinel from configs import dify_config -class RedisClientWrapper(redis.Redis): +class RedisClientWrapper: """ A wrapper class for the Redis client that addresses the issue where the global `redis_client` variable cannot be updated when a new Redis instance is returned @@ -71,6 +72,12 @@ def init_app(app): ) master = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **redis_params) redis_client.initialize(master) + elif dify_config.REDIS_USE_CLUSTERS: + nodes = [ + ClusterNode(host=node.split(":")[0], port=int(node.split.split(":")[1])) + for node in dify_config.REDIS_CLUSTERS.split(",") + ] + redis_client.initialize(RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD)) else: redis_params.update( { diff --git a/api/tests/unit_tests/core/test_model_manager.py b/api/tests/unit_tests/core/test_model_manager.py index 2808b5b0fa..d98e9f6bad 100644 --- a/api/tests/unit_tests/core/test_model_manager.py +++ b/api/tests/unit_tests/core/test_model_manager.py @@ -1,10 +1,12 @@ -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest +import redis from core.entities.provider_entities import ModelLoadBalancingConfiguration from core.model_manager import LBModelManager from core.model_runtime.entities.model_entities import ModelType +from extensions.ext_redis import redis_client @pytest.fixture @@ -38,6 +40,9 @@ def lb_model_manager(): def test_lb_model_manager_fetch_next(mocker, lb_model_manager): + # initialize redis client + redis_client.initialize(redis.Redis()) + assert len(lb_model_manager._load_balancing_configs) == 3 config1 = lb_model_manager._load_balancing_configs[0] @@ -55,12 +60,13 @@ def test_lb_model_manager_fetch_next(mocker, lb_model_manager): start_index += 1 return start_index - mocker.patch("redis.Redis.incr", side_effect=incr) - mocker.patch("redis.Redis.set", return_value=None) - mocker.patch("redis.Redis.expire", return_value=None) + with ( + patch.object(redis_client, "incr", side_effect=incr), + patch.object(redis_client, "set", return_value=None), + patch.object(redis_client, "expire", return_value=None), + ): + config = lb_model_manager.fetch_next() + assert config == config2 - config = lb_model_manager.fetch_next() - assert config == config2 - - config = lb_model_manager.fetch_next() - assert config == config3 + config = lb_model_manager.fetch_next() + assert config == config3 diff --git a/docker/.env.example b/docker/.env.example index be8d72339f..d29c66535d 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -240,6 +240,12 @@ REDIS_SENTINEL_USERNAME= REDIS_SENTINEL_PASSWORD= REDIS_SENTINEL_SOCKET_TIMEOUT=0.1 +# List of Redis Cluster nodes. If Cluster mode is enabled, provide at least one Cluster IP and port. +# Format: `:,:,:` +REDIS_USE_CLUSTERS=false +REDIS_CLUSTERS= +REDIS_CLUSTERS_PASSWORD= + # ------------------------------ # Celery Configuration # ------------------------------ diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index b6caff90d9..06b04d90b6 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -55,6 +55,9 @@ x-shared-env: &shared-api-worker-env REDIS_SENTINEL_USERNAME: ${REDIS_SENTINEL_USERNAME:-} REDIS_SENTINEL_PASSWORD: ${REDIS_SENTINEL_PASSWORD:-} REDIS_SENTINEL_SOCKET_TIMEOUT: ${REDIS_SENTINEL_SOCKET_TIMEOUT:-0.1} + REDIS_CLUSTERS: ${REDIS_CLUSTERS:-} + REDIS_USE_CLUSTERS: ${REDIS_USE_CLUSTERS:-false} + REDIS_CLUSTERS_PASSWORD: ${REDIS_CLUSTERS_PASSWORD:-} ACCESS_TOKEN_EXPIRE_MINUTES: ${ACCESS_TOKEN_EXPIRE_MINUTES:-60} CELERY_BROKER_URL: ${CELERY_BROKER_URL:-redis://:difyai123456@redis:6379/1} BROKER_USE_SSL: ${BROKER_USE_SSL:-false}