Perf/indexing status page (onyx-dot-app#5142)

* indexing status optimization first draft

* refactor: update pagination logic and enhance UI for indexing status table

* add index attempt pruning job and display federated connectors in index status page

* update celery worker command to include index_attempt_cleanup queue

* refactor: enhance indexing status table and remove deprecated components

* mypy fix

* address review comments

* fix pagination reset issue

* add TODO for optimizing connector materialization and performance in future deployments

* enhance connector indexing status retrieval by adding 'get_all_connectors' option and updating pagination logic

* refactor: transition to paginated connector indexing status retrieval and update related components

* fix: initialize latest_index_attempt_docs_indexed to 0 in CCPairIndexingStatusTable component

* feat: add mock connector file support for indexing status retrieval and update indexing_statuses type to Sequence

* mypy fix

* refactor: rename indexing status endpoint to simplify API and update related components

Author: Subash-Mohan

.vscode/launch.template.jsonc

@@ -192,7 +192,7 @@
           "--loglevel=INFO",
           "--hostname=light@%n",
           "-Q",
-          "vespa_metadata_sync,connector_deletion,doc_permissions_upsert"
+          "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,index_attempt_cleanup"
         ],
         "presentation": {
           "group": "2"

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

@@ -62,6 +62,15 @@
             "expires": BEAT_EXPIRES_DEFAULT,
         },
     },
+    {
+        "name": "check-for-index-attempt-cleanup",
+        "task": OnyxCeleryTask.CHECK_FOR_INDEX_ATTEMPT_CLEANUP,
+        "schedule": timedelta(hours=1),
+        "options": {
+            "priority": OnyxCeleryPriority.LOW,
+            "expires": BEAT_EXPIRES_DEFAULT,
+        },
+    },
     {
         "name": "check-for-connector-deletion",
         "task": OnyxCeleryTask.CHECK_FOR_CONNECTOR_DELETION,

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

@@ -32,6 +32,8 @@
 from onyx.background.indexing.checkpointing_utils import (
     get_index_attempts_with_old_checkpoints,
 )
+from onyx.background.indexing.index_attempt_utils import cleanup_index_attempts
+from onyx.background.indexing.index_attempt_utils import get_old_index_attempts
 from onyx.configs.app_configs import MANAGED_VESPA
 from onyx.configs.app_configs import VESPA_CLOUD_CERT_PATH
 from onyx.configs.app_configs import VESPA_CLOUD_KEY_PATH
@@ -109,6 +111,7 @@
 # Heartbeat timeout: if no heartbeat received for 30 minutes, consider it dead
 # This should be much longer than INDEXING_WORKER_HEARTBEAT_INTERVAL (30s)
 HEARTBEAT_TIMEOUT_SECONDS = 30 * 60  # 30 minutes
+INDEX_ATTEMPT_BATCH_SIZE = 500
 
 
 def _get_fence_validation_block_expiration() -> int:
@@ -987,6 +990,95 @@ def cleanup_checkpoint_task(
         )
 
 
+# primary
+@shared_task(
+    name=OnyxCeleryTask.CHECK_FOR_INDEX_ATTEMPT_CLEANUP,
+    soft_time_limit=300,
+    bind=True,
+)
+def check_for_index_attempt_cleanup(self: Task, *, tenant_id: str) -> None:
+    """Clean up old index attempts that are older than 7 days."""
+    locked = False
+    redis_client = get_redis_client(tenant_id=tenant_id)
+    lock: RedisLock = redis_client.lock(
+        OnyxRedisLocks.CHECK_INDEX_ATTEMPT_CLEANUP_BEAT_LOCK,
+        timeout=CELERY_GENERIC_BEAT_LOCK_TIMEOUT,
+    )
+
+    # these tasks should never overlap
+    if not lock.acquire(blocking=False):
+        task_logger.info(
+            f"check_for_index_attempt_cleanup - Lock not acquired: tenant={tenant_id}"
+        )
+        return None
+
+    try:
+        locked = True
+        batch_size = INDEX_ATTEMPT_BATCH_SIZE
+        with get_session_with_current_tenant() as db_session:
+            old_attempts = get_old_index_attempts(db_session)
+            # We need to batch this because during the initial run, the system might have a large number
+            # of index attempts since they were never deleted. After that, the number will be
+            # significantly lower.
+            if len(old_attempts) == 0:
+                task_logger.info(
+                    "check_for_index_attempt_cleanup - No index attempts to cleanup"
+                )
+                return
+
+            for i in range(0, len(old_attempts), batch_size):
+                batch = old_attempts[i : i + batch_size]
+                task_logger.info(
+                    f"check_for_index_attempt_cleanup - Cleaning up index attempts {len(batch)}"
+                )
+                self.app.send_task(
+                    OnyxCeleryTask.CLEANUP_INDEX_ATTEMPT,
+                    kwargs={
+                        "index_attempt_ids": [attempt.id for attempt in batch],
+                        "tenant_id": tenant_id,
+                    },
+                    queue=OnyxCeleryQueues.INDEX_ATTEMPT_CLEANUP,
+                    priority=OnyxCeleryPriority.MEDIUM,
+                )
+    except Exception:
+        task_logger.exception("Unexpected exception during index attempt cleanup check")
+        return None
+    finally:
+        if locked:
+            if lock.owned():
+                lock.release()
+            else:
+                task_logger.error(
+                    "check_for_index_attempt_cleanup - Lock not owned on completion: "
+                    f"tenant={tenant_id}"
+                )
+
+
+# light worker
+@shared_task(
+    name=OnyxCeleryTask.CLEANUP_INDEX_ATTEMPT,
+    bind=True,
+)
+def cleanup_index_attempt_task(
+    self: Task, *, index_attempt_ids: list[int], tenant_id: str
+) -> None:
+    """Clean up an index attempt"""
+    start = time.monotonic()
+
+    try:
+        with get_session_with_current_tenant() as db_session:
+            cleanup_index_attempts(db_session, index_attempt_ids)
+
+    finally:
+        elapsed = time.monotonic() - start
+
+        task_logger.info(
+            f"cleanup_index_attempt_task completed: tenant_id={tenant_id} "
+            f"index_attempt_ids={index_attempt_ids} "
+            f"elapsed={elapsed:.2f}"
+        )
+
+
 class DocumentProcessingBatch(BaseModel):
     """Data structure for a document processing batch."""
 

backend/onyx/background/indexing/checkpointing_utils.py

@@ -6,6 +6,7 @@
 from sqlalchemy.orm import Session
 
 from onyx.configs.constants import FileOrigin
+from onyx.configs.constants import NUM_DAYS_TO_KEEP_CHECKPOINTS
 from onyx.connectors.interfaces import BaseConnector
 from onyx.connectors.interfaces import CheckpointedConnector
 from onyx.connectors.models import ConnectorCheckpoint
@@ -163,16 +164,16 @@ def get_latest_valid_checkpoint(
 
 
 def get_index_attempts_with_old_checkpoints(
-    db_session: Session, days_to_keep: int = 7
+    db_session: Session, days_to_keep: int = NUM_DAYS_TO_KEEP_CHECKPOINTS
 ) -> list[IndexAttempt]:
     """Get all index attempts with checkpoints older than the specified number of days.
 
     Args:
         db_session: The database session
-        days_to_keep: Number of days to keep checkpoints for (default: 7)
+        days_to_keep: Number of days to keep checkpoints for (default: NUM_DAYS_TO_KEEP_CHECKPOINTS)
 
     Returns:
-        Number of checkpoints deleted
+        List of IndexAttempt objects with old checkpoints
     """
     cutoff_date = get_db_current_time(db_session) - timedelta(days=days_to_keep)
 

backend/onyx/background/indexing/index_attempt_utils.py

@@ -0,0 +1,27 @@
+from datetime import timedelta
+
+from sqlalchemy.orm import Session
+
+from onyx.configs.constants import NUM_DAYS_TO_KEEP_INDEX_ATTEMPTS
+from onyx.db.engine.time_utils import get_db_current_time
+from onyx.db.models import IndexAttempt
+
+
+def get_old_index_attempts(
+    db_session: Session, days_to_keep: int = NUM_DAYS_TO_KEEP_INDEX_ATTEMPTS
+) -> list[IndexAttempt]:
+    """Get all index attempts older than the specified number of days."""
+    cutoff_date = get_db_current_time(db_session) - timedelta(days=days_to_keep)
+    return (
+        db_session.query(IndexAttempt)
+        .filter(IndexAttempt.time_created < cutoff_date)
+        .all()
+    )
+
+
+def cleanup_index_attempts(db_session: Session, index_attempt_ids: list[int]) -> None:
+    """Clean up multiple index attempts"""
+    db_session.query(IndexAttempt).filter(
+        IndexAttempt.id.in_(index_attempt_ids)
+    ).delete(synchronize_session=False)
+    db_session.commit()

backend/onyx/configs/constants.py

@@ -326,7 +326,7 @@ class OnyxCeleryQueues:
     CONNECTOR_DELETION = "connector_deletion"
     LLM_MODEL_UPDATE = "llm_model_update"
     CHECKPOINT_CLEANUP = "checkpoint_cleanup"
-
+    INDEX_ATTEMPT_CLEANUP = "index_attempt_cleanup"
     # Heavy queue
     CONNECTOR_PRUNING = "connector_pruning"
     CONNECTOR_DOC_PERMISSIONS_SYNC = "connector_doc_permissions_sync"
@@ -354,6 +354,7 @@ class OnyxRedisLocks:
     CHECK_PRUNE_BEAT_LOCK = "da_lock:check_prune_beat"
     CHECK_INDEXING_BEAT_LOCK = "da_lock:check_indexing_beat"
     CHECK_CHECKPOINT_CLEANUP_BEAT_LOCK = "da_lock:check_checkpoint_cleanup_beat"
+    CHECK_INDEX_ATTEMPT_CLEANUP_BEAT_LOCK = "da_lock:check_index_attempt_cleanup_beat"
     CHECK_CONNECTOR_DOC_PERMISSIONS_SYNC_BEAT_LOCK = (
         "da_lock:check_connector_doc_permissions_sync_beat"
     )
@@ -455,6 +456,10 @@ class OnyxCeleryTask:
     CHECK_FOR_CHECKPOINT_CLEANUP = "check_for_checkpoint_cleanup"
     CLEANUP_CHECKPOINT = "cleanup_checkpoint"
 
+    # Connector index attempt cleanup
+    CHECK_FOR_INDEX_ATTEMPT_CLEANUP = "check_for_index_attempt_cleanup"
+    CLEANUP_INDEX_ATTEMPT = "cleanup_index_attempt"
+
     MONITOR_BACKGROUND_PROCESSES = "monitor_background_processes"
     MONITOR_CELERY_QUEUES = "monitor_celery_queues"
     MONITOR_PROCESS_MEMORY = "monitor_process_memory"
@@ -512,3 +517,8 @@ class OnyxCeleryTask:
 class OnyxCallTypes(str, Enum):
     FIREFLIES = "FIREFLIES"
     GONG = "GONG"
+
+
+NUM_DAYS_TO_KEEP_CHECKPOINTS = 7
+# checkpoints are queried based on index attempts, so we need to keep index attempts for one more day
+NUM_DAYS_TO_KEEP_INDEX_ATTEMPTS = NUM_DAYS_TO_KEEP_CHECKPOINTS + 1

backend/onyx/db/connector_credential_pair.py

@@ -116,12 +116,13 @@ def get_connector_credential_pairs_for_user(
     eager_load_credential: bool = False,
     eager_load_user: bool = False,
     include_user_files: bool = False,
+    order_by_desc: bool = False,
+    source: DocumentSource | None = None,
 ) -> list[ConnectorCredentialPair]:
     if eager_load_user:
         assert (
             eager_load_credential
         ), "eager_load_credential must be True if eager_load_user is True"
-
     stmt = select(ConnectorCredentialPair).distinct()
 
     if eager_load_connector:
@@ -134,12 +135,21 @@ def get_connector_credential_pairs_for_user(
         stmt = stmt.options(load_opts)
 
     stmt = _add_user_filters(stmt, user, get_editable)
+
+    if source:
+        stmt = stmt.join(ConnectorCredentialPair.connector).where(
+            Connector.source == source.value
+        )
+
     if ids:
         stmt = stmt.where(ConnectorCredentialPair.id.in_(ids))
 
     if not include_user_files:
         stmt = stmt.where(ConnectorCredentialPair.is_user_file != True)  # noqa: E712
 
+    if order_by_desc:
+        stmt = stmt.order_by(desc(ConnectorCredentialPair.id))
+
     return list(db_session.scalars(stmt).unique().all())
 
 
@@ -153,16 +163,20 @@ def get_connector_credential_pairs_for_user_parallel(
     eager_load_connector: bool = False,
     eager_load_credential: bool = False,
     eager_load_user: bool = False,
+    order_by_desc: bool = False,
+    source: DocumentSource | None = None,
 ) -> list[ConnectorCredentialPair]:
     with get_session_with_current_tenant() as db_session:
         return get_connector_credential_pairs_for_user(
-            db_session,
-            user,
-            get_editable,
-            ids,
-            eager_load_connector,
-            eager_load_credential,
-            eager_load_user,
+            db_session=db_session,
+            user=user,
+            get_editable=get_editable,
+            ids=ids,
+            eager_load_connector=eager_load_connector,
+            eager_load_credential=eager_load_credential,
+            eager_load_user=eager_load_user,
+            order_by_desc=order_by_desc,
+            source=source,
         )
 
 

backend/onyx/db/federated.py

@@ -8,6 +8,7 @@
 from sqlalchemy.orm import Session
 
 from onyx.configs.constants import FederatedConnectorSource
+from onyx.db.engine.sql_engine import get_session_with_current_tenant
 from onyx.db.models import DocumentSet
 from onyx.db.models import FederatedConnector
 from onyx.db.models import FederatedConnector__DocumentSet
@@ -39,6 +40,11 @@ def fetch_all_federated_connectors(db_session: Session) -> list[FederatedConnect
     return list(result.scalars().all())
 
 
+def fetch_all_federated_connectors_parallel() -> list[FederatedConnector]:
+    with get_session_with_current_tenant() as db_session:
+        return fetch_all_federated_connectors(db_session)
+
+
 def validate_federated_connector_credentials(
     source: FederatedConnectorSource,
     credentials: dict[str, Any],