use a larger scan_iter value for performance

backend/onyx/background/celery/tasks/indexing/tasks.py

@@ -63,6 +63,7 @@
 from onyx.redis.redis_connector_index import RedisConnectorIndexPayload
 from onyx.redis.redis_pool import get_redis_client
 from onyx.redis.redis_pool import redis_lock_dump
+from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT
 from onyx.utils.logger import setup_logger
 from onyx.utils.variable_functionality import global_version
 from shared_configs.configs import INDEXING_MODEL_SERVER_HOST
@@ -262,7 +263,7 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
                 tenant_id == "tenant_i-043470d740845ec56"
                 or tenant_id == "tenant_82b497ce-88aa-4fbd-841a-92cae43529c8"
             ):
-                logger.info(
+                task_logger.info(
                     f"check_for_indexing lock: "
                     f"tenant={tenant_id} "
                     f"cc_pair={cc_pair_id} "
@@ -417,7 +418,9 @@ def validate_indexing_fences(
     )
 
     # validate all existing indexing jobs
-    for key_bytes in r.scan_iter(RedisConnectorIndex.FENCE_PREFIX + "*"):
+    for key_bytes in r.scan_iter(
+        RedisConnectorIndex.FENCE_PREFIX + "*", count=SCAN_ITER_COUNT_DEFAULT
+    ):
         lock_beat.reacquire()
         with get_session_with_tenant(tenant_id) as db_session:
             validate_indexing_fence(

backend/onyx/background/celery/tasks/vespa/tasks.py

@@ -71,6 +71,7 @@
 from onyx.redis.redis_document_set import RedisDocumentSet
 from onyx.redis.redis_pool import get_redis_client
 from onyx.redis.redis_pool import redis_lock_dump
+from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT
 from onyx.redis.redis_usergroup import RedisUserGroup
 from onyx.utils.logger import setup_logger
 from onyx.utils.variable_functionality import fetch_versioned_implementation
@@ -840,7 +841,9 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool | None:
 
         phase_start = time.monotonic()
         lock_beat.reacquire()
-        for key_bytes in r.scan_iter(RedisConnectorDelete.FENCE_PREFIX + "*"):
+        for key_bytes in r.scan_iter(
+            RedisConnectorDelete.FENCE_PREFIX + "*", count=SCAN_ITER_COUNT_DEFAULT
+        ):
             monitor_connector_deletion_taskset(tenant_id, key_bytes, r)
             lock_beat.reacquire()
 
@@ -851,7 +854,9 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool | None:
 
         phase_start = time.monotonic()
         lock_beat.reacquire()
-        for key_bytes in r.scan_iter(RedisDocumentSet.FENCE_PREFIX + "*"):
+        for key_bytes in r.scan_iter(
+            RedisDocumentSet.FENCE_PREFIX + "*", count=SCAN_ITER_COUNT_DEFAULT
+        ):
             with get_session_with_tenant(tenant_id) as db_session:
                 monitor_document_set_taskset(tenant_id, key_bytes, r, db_session)
             lock_beat.reacquire()
@@ -860,7 +865,9 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool | None:
 
         phase_start = time.monotonic()
         lock_beat.reacquire()
-        for key_bytes in r.scan_iter(RedisUserGroup.FENCE_PREFIX + "*"):
+        for key_bytes in r.scan_iter(
+            RedisUserGroup.FENCE_PREFIX + "*", count=SCAN_ITER_COUNT_DEFAULT
+        ):
             monitor_usergroup_taskset = fetch_versioned_implementation_with_fallback(
                 "onyx.background.celery.tasks.vespa.tasks",
                 "monitor_usergroup_taskset",
@@ -874,7 +881,9 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool | None:
 
         phase_start = time.monotonic()
         lock_beat.reacquire()
-        for key_bytes in r.scan_iter(RedisConnectorPrune.FENCE_PREFIX + "*"):
+        for key_bytes in r.scan_iter(
+            RedisConnectorPrune.FENCE_PREFIX + "*", count=SCAN_ITER_COUNT_DEFAULT
+        ):
             with get_session_with_tenant(tenant_id) as db_session:
                 monitor_ccpair_pruning_taskset(tenant_id, key_bytes, r, db_session)
             lock_beat.reacquire()
@@ -883,7 +892,9 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool | None:
 
         phase_start = time.monotonic()
         lock_beat.reacquire()
-        for key_bytes in r.scan_iter(RedisConnectorIndex.FENCE_PREFIX + "*"):
+        for key_bytes in r.scan_iter(
+            RedisConnectorIndex.FENCE_PREFIX + "*", count=SCAN_ITER_COUNT_DEFAULT
+        ):
             with get_session_with_tenant(tenant_id) as db_session:
                 monitor_ccpair_indexing_taskset(tenant_id, key_bytes, r, db_session)
             lock_beat.reacquire()
@@ -892,7 +903,10 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool | None:
 
         phase_start = time.monotonic()
         lock_beat.reacquire()
-        for key_bytes in r.scan_iter(RedisConnectorPermissionSync.FENCE_PREFIX + "*"):
+        for key_bytes in r.scan_iter(
+            RedisConnectorPermissionSync.FENCE_PREFIX + "*",
+            count=SCAN_ITER_COUNT_DEFAULT,
+        ):
             with get_session_with_tenant(tenant_id) as db_session:
                 monitor_ccpair_permissions_taskset(tenant_id, key_bytes, r, db_session)
             lock_beat.reacquire()

backend/onyx/redis/redis_connector_doc_perm_sync.py

@@ -13,6 +13,7 @@
 from onyx.configs.constants import OnyxCeleryPriority
 from onyx.configs.constants import OnyxCeleryQueues
 from onyx.configs.constants import OnyxCeleryTask
+from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT
 
 
 class RedisConnectorPermissionSyncPayload(BaseModel):
@@ -68,7 +69,10 @@ def get_remaining(self) -> int:
     def get_active_task_count(self) -> int:
         """Count of active permission sync tasks"""
         count = 0
-        for _ in self.redis.scan_iter(RedisConnectorPermissionSync.FENCE_PREFIX + "*"):
+        for _ in self.redis.scan_iter(
+            RedisConnectorPermissionSync.FENCE_PREFIX + "*",
+            count=SCAN_ITER_COUNT_DEFAULT,
+        ):
             count += 1
         return count
 

backend/onyx/redis/redis_connector_ext_group_sync.py

@@ -7,6 +7,8 @@
 from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 
+from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT
+
 
 class RedisConnectorExternalGroupSyncPayload(BaseModel):
     started: datetime | None
@@ -63,7 +65,8 @@ def get_active_task_count(self) -> int:
         """Count of active external group syncing tasks"""
         count = 0
         for _ in self.redis.scan_iter(
-            RedisConnectorExternalGroupSync.FENCE_PREFIX + "*"
+            RedisConnectorExternalGroupSync.FENCE_PREFIX + "*",
+            count=SCAN_ITER_COUNT_DEFAULT,
         ):
             count += 1
         return count

backend/onyx/redis/redis_connector_prune.py

@@ -12,6 +12,7 @@
 from onyx.configs.constants import OnyxCeleryQueues
 from onyx.configs.constants import OnyxCeleryTask
 from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
+from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT
 
 
 class RedisConnectorPrune:
@@ -63,7 +64,9 @@ def get_remaining(self) -> int:
     def get_active_task_count(self) -> int:
         """Count of active pruning tasks"""
         count = 0
-        for key in self.redis.scan_iter(RedisConnectorPrune.FENCE_PREFIX + "*"):
+        for key in self.redis.scan_iter(
+            RedisConnectorPrune.FENCE_PREFIX + "*", count=SCAN_ITER_COUNT_DEFAULT
+        ):
             count += 1
         return count
 

backend/onyx/redis/redis_pool.py

@@ -28,6 +28,8 @@
 
 logger = setup_logger()
 
+SCAN_ITER_COUNT_DEFAULT = 4096
+
 
 class TenantRedis(redis.Redis):
     def __init__(self, tenant_id: str, *args: Any, **kwargs: Any) -> None: