Browse Source

refactor: change MAX_RETRY_COUNT as env

Signed-off-by: Sihyeon Jang <sihyeon.jang@navercorp.com>
Sihyeon Jang 2 months ago
parent
commit
65882c30cb
2 changed files with 14 additions and 6 deletions
  1. 9 0
      backend/open_webui/env.py
  2. 5 6
      backend/open_webui/utils/redis.py

+ 9 - 0
backend/open_webui/env.py

@@ -349,6 +349,15 @@ REDIS_KEY_PREFIX = os.environ.get("REDIS_KEY_PREFIX", "open-webui")
 REDIS_SENTINEL_HOSTS = os.environ.get("REDIS_SENTINEL_HOSTS", "")
 REDIS_SENTINEL_PORT = os.environ.get("REDIS_SENTINEL_PORT", "26379")
 
+# Maximum number of retries for Redis operations when using Sentinel fail-over
+REDIS_SENTINEL_MAX_RETRY_COUNT = os.environ.get("REDIS_SENTINEL_MAX_RETRY_COUNT", "2")
+try:
+    REDIS_SENTINEL_MAX_RETRY_COUNT = int(REDIS_SENTINEL_MAX_RETRY_COUNT)
+    if REDIS_SENTINEL_MAX_RETRY_COUNT < 1:
+        REDIS_SENTINEL_MAX_RETRY_COUNT = 2
+except ValueError:
+    REDIS_SENTINEL_MAX_RETRY_COUNT = 2
+
 ####################################
 # UVICORN WORKERS
 ####################################

+ 5 - 6
backend/open_webui/utils/redis.py

@@ -3,8 +3,7 @@ from urllib.parse import urlparse
 
 import redis
 
-
-MAX_RETRY_COUNT = 2
+from open_webui.env import REDIS_SENTINEL_MAX_RETRY_COUNT
 
 
 class SentinelRedisProxy:
@@ -31,7 +30,7 @@ class SentinelRedisProxy:
         if self._async_mode:
 
             async def _wrapped(*args, **kwargs):
-                for i in range(MAX_RETRY_COUNT):
+                for i in range(REDIS_SENTINEL_MAX_RETRY_COUNT):
                     try:
                         method = getattr(self._master(), item)
                         result = method(*args, **kwargs)
@@ -42,7 +41,7 @@ class SentinelRedisProxy:
                         redis.exceptions.ConnectionError,
                         redis.exceptions.ReadOnlyError,
                     ) as e:
-                        if i < MAX_RETRY_COUNT - 1:
+                        if i < REDIS_SENTINEL_MAX_RETRY_COUNT - 1:
                             continue
                         raise e from e
 
@@ -51,7 +50,7 @@ class SentinelRedisProxy:
         else:
 
             def _wrapped(*args, **kwargs):
-                for i in range(MAX_RETRY_COUNT):
+                for i in range(REDIS_SENTINEL_MAX_RETRY_COUNT):
                     try:
                         method = getattr(self._master(), item)
                         return method(*args, **kwargs)
@@ -59,7 +58,7 @@ class SentinelRedisProxy:
                         redis.exceptions.ConnectionError,
                         redis.exceptions.ReadOnlyError,
                     ) as e:
-                        if i < MAX_RETRY_COUNT - 1:
+                        if i < REDIS_SENTINEL_MAX_RETRY_COUNT - 1:
                             continue
                         raise e from e