perf(celery): cut worker boot memory by ~28% with lazy imports (#14632)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

backend/ee/onyx/background/celery/tasks/query_history/tasks.py

@@ -4,10 +4,6 @@ from datetime import datetime
 
 from celery import Task, shared_task
 
-from ee.onyx.server.query_history.api import (
-    ONYX_ANONYMIZED_EMAIL,
-    fetch_and_process_chat_session_history,
-)
 from ee.onyx.server.query_history.models import QuestionAnswerPairSnapshot
 from onyx.background.task_utils import construct_query_history_report_name
 from onyx.configs.app_configs import JOB_TIMEOUT
@@ -47,6 +43,11 @@ def export_query_history_task(
     # Need to include the tenant_id since the TenantAwareTask needs this
     tenant_id: str,  # noqa: ARG001
 ) -> None:
+    from ee.onyx.server.query_history.api import (
+        ONYX_ANONYMIZED_EMAIL,
+        fetch_and_process_chat_session_history,
+    )
+
     if not self.request.id:
         raise RuntimeError("No task id defined for this task; cannot identify it")
 

backend/ee/onyx/background/celery/tasks/sso_domain_revalidation/tasks.py

@@ -7,7 +7,6 @@ domain whose DNS TXT proof no longer resolves.
 
 from celery import shared_task
 
-from ee.onyx.auth.sso_domain_verification import revalidate_tenant_domains
 from ee.onyx.db.tenant_sso_domain import reproject_tenant_login_domains
 from onyx.configs.app_configs import JOB_TIMEOUT
 from onyx.configs.constants import OnyxCeleryTask
@@ -26,6 +25,8 @@ def revalidate_sso_domains_task(*, tenant_id: str) -> None:
     """Fanned out per tenant by cloud_beat_task_generator. The re-projection is
     isolated so its failure can't skip the DNS re-check, which is the part that
     drops routing for a domain whose proof is gone."""
+    from ee.onyx.auth.sso_domain_verification import revalidate_tenant_domains
+
     try:
         reproject_tenant_login_domains(tenant_id)
     except Exception:

backend/ee/onyx/background/celery/tasks/tenant_provisioning/tasks.py

@@ -9,12 +9,6 @@ import uuid
 from celery import Task, shared_task
 from redis.lock import Lock as RedisLock
 
-from ee.onyx.server.tenants.provisioning import setup_tenant
-from ee.onyx.server.tenants.schema_management import (
-    create_schema_if_not_exists,
-    get_current_alembic_version,
-    run_alembic_migrations,
-)
 from onyx.background.celery.apps.app_base import task_logger
 from onyx.configs.app_configs import TARGET_AVAILABLE_TENANTS
 from onyx.configs.constants import (
@@ -135,6 +129,11 @@ def _migrate_stale_pool_tenants() -> None:
     tenants are always current so that signup doesn't hit schema mismatches
     (e.g. missing columns added after the tenant was pre-provisioned).
     """
+    from ee.onyx.server.tenants.schema_management import (
+        get_current_alembic_version,
+        run_alembic_migrations,
+    )
+
     with get_session_with_shared_schema() as db_session:
         pool_tenants = db_session.query(AvailableTenant).all()
         tenant_ids = [t.tenant_id for t in pool_tenants]
@@ -178,6 +177,13 @@ def pre_provision_tenant() -> bool:
     """
     # The MULTI_TENANT check is now done at the caller level (check_available_tenants)
     # rather than inside this function
+    # Imported here: provisioning reaches every tool implementation (~75 MB) and
+    # only multi-tenant deployments run this.
+    from ee.onyx.server.tenants.provisioning import setup_tenant
+    from ee.onyx.server.tenants.schema_management import (
+        create_schema_if_not_exists,
+        get_current_alembic_version,
+    )
 
     r = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)
     lock_provision: RedisLock = r.lock(

backend/ee/onyx/background/celery/tasks/usage_reporting/tasks.py

@@ -3,7 +3,6 @@ from uuid import UUID
 
 from celery import Task, shared_task
 
-from ee.onyx.server.reporting.usage_export_generation import create_new_usage_report
 from onyx.configs.app_configs import JOB_TIMEOUT
 from onyx.configs.constants import OnyxCeleryTask
 from onyx.db.engine.sql_engine import get_session_with_current_tenant
@@ -29,6 +28,8 @@ def generate_usage_report_task(
     report_id: str | None = None,
 ) -> None:
     """User-initiated usage report generation task"""
+    from ee.onyx.server.reporting.usage_export_generation import create_new_usage_report
+
     # Parse period if provided
     period = None
     if period_from and period_to:

backend/ee/onyx/external_permissions/sync_params.py

@@ -1,4 +1,4 @@
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 from typing import TYPE_CHECKING, Optional
 
 from pydantic import BaseModel
@@ -21,19 +21,6 @@ from ee.onyx.configs.app_configs import (
     SLACK_PERMISSION_DOC_SYNC_FREQUENCY,
     TEAMS_PERMISSION_DOC_SYNC_FREQUENCY,
 )
-from ee.onyx.external_permissions.box.doc_sync import box_doc_sync
-from ee.onyx.external_permissions.box.group_sync import box_group_sync
-from ee.onyx.external_permissions.canvas.doc_sync import canvas_doc_sync
-from ee.onyx.external_permissions.canvas.group_sync import canvas_group_sync
-from ee.onyx.external_permissions.confluence.doc_sync import confluence_doc_sync
-from ee.onyx.external_permissions.confluence.group_sync import confluence_group_sync
-from ee.onyx.external_permissions.github.doc_sync import github_doc_sync
-from ee.onyx.external_permissions.github.group_sync import github_group_sync
-from ee.onyx.external_permissions.gmail.doc_sync import gmail_doc_sync
-from ee.onyx.external_permissions.google_drive.doc_sync import gdrive_doc_sync
-from ee.onyx.external_permissions.google_drive.group_sync import gdrive_group_sync
-from ee.onyx.external_permissions.jira.doc_sync import jira_doc_sync
-from ee.onyx.external_permissions.jira.group_sync import jira_group_sync
 from ee.onyx.external_permissions.perm_sync_types import (
     CensoringFuncType,
     DocSyncFuncType,
@@ -41,21 +28,157 @@ from ee.onyx.external_permissions.perm_sync_types import (
     FetchAllDocumentsIdsFunction,
     GroupSyncFuncType,
 )
-from ee.onyx.external_permissions.salesforce.postprocessing import (
-    censor_salesforce_chunks,
-)
-from ee.onyx.external_permissions.sharepoint.doc_sync import sharepoint_doc_sync
-from ee.onyx.external_permissions.sharepoint.group_sync import sharepoint_group_sync
-from ee.onyx.external_permissions.slack.doc_sync import slack_doc_sync
-from ee.onyx.external_permissions.teams.doc_sync import teams_doc_sync
 from onyx.configs.constants import DocumentSource
 
 if TYPE_CHECKING:
-    from onyx.access.models import DocExternalAccess  # noqa
+    from ee.onyx.db.external_perm import ExternalUserGroup  # noqa
+    from onyx.access.models import DocExternalAccess, ElementExternalAccess  # noqa
+    from onyx.context.search.models import InferenceChunk  # noqa
     from onyx.db.models import ConnectorCredentialPair  # noqa
     from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface  # noqa
 
 
+# Sync implementations import their connector SDKs (box_sdk_gen, office365,
+# google, github, ...), which cost ~60 MB. Load them on first call so workers
+# that only read this config stay small.
+def _lazy_doc_sync(load: Callable[[], DocSyncFuncType]) -> DocSyncFuncType:
+    def run(
+        cc_pair: "ConnectorCredentialPair",
+        fetch_all_docs_fn: FetchAllDocumentsFunction,
+        fetch_all_docs_ids_fn: FetchAllDocumentsIdsFunction,
+        callback: Optional["IndexingHeartbeatInterface"],
+    ) -> Generator["ElementExternalAccess", None, None]:
+        return load()(cc_pair, fetch_all_docs_fn, fetch_all_docs_ids_fn, callback)
+
+    return run
+
+
+def _lazy_group_sync(load: Callable[[], GroupSyncFuncType]) -> GroupSyncFuncType:
+    def run(
+        tenant_id: str, cc_pair: "ConnectorCredentialPair"
+    ) -> Generator["ExternalUserGroup", None, None]:
+        return load()(tenant_id, cc_pair)
+
+    return run
+
+
+def _lazy_censoring(load: Callable[[], CensoringFuncType]) -> CensoringFuncType:
+    def run(chunks: list["InferenceChunk"], user_email: str) -> list["InferenceChunk"]:
+        return load()(chunks, user_email)
+
+    return run
+
+
+def _load_box_doc_sync() -> DocSyncFuncType:
+    from ee.onyx.external_permissions.box.doc_sync import box_doc_sync
+
+    return box_doc_sync
+
+
+def _load_box_group_sync() -> GroupSyncFuncType:
+    from ee.onyx.external_permissions.box.group_sync import box_group_sync
+
+    return box_group_sync
+
+
+def _load_canvas_doc_sync() -> DocSyncFuncType:
+    from ee.onyx.external_permissions.canvas.doc_sync import canvas_doc_sync
+
+    return canvas_doc_sync
+
+
+def _load_canvas_group_sync() -> GroupSyncFuncType:
+    from ee.onyx.external_permissions.canvas.group_sync import canvas_group_sync
+
+    return canvas_group_sync
+
+
+def _load_confluence_doc_sync() -> DocSyncFuncType:
+    from ee.onyx.external_permissions.confluence.doc_sync import confluence_doc_sync
+
+    return confluence_doc_sync
+
+
+def _load_confluence_group_sync() -> GroupSyncFuncType:
+    from ee.onyx.external_permissions.confluence.group_sync import confluence_group_sync
+
+    return confluence_group_sync
+
+
+def _load_github_doc_sync() -> DocSyncFuncType:
+    from ee.onyx.external_permissions.github.doc_sync import github_doc_sync
+
+    return github_doc_sync
+
+
+def _load_github_group_sync() -> GroupSyncFuncType:
+    from ee.onyx.external_permissions.github.group_sync import github_group_sync
+
+    return github_group_sync
+
+
+def _load_gmail_doc_sync() -> DocSyncFuncType:
+    from ee.onyx.external_permissions.gmail.doc_sync import gmail_doc_sync
+
+    return gmail_doc_sync
+
+
+def _load_gdrive_doc_sync() -> DocSyncFuncType:
+    from ee.onyx.external_permissions.google_drive.doc_sync import gdrive_doc_sync
+
+    return gdrive_doc_sync
+
+
+def _load_gdrive_group_sync() -> GroupSyncFuncType:
+    from ee.onyx.external_permissions.google_drive.group_sync import gdrive_group_sync
+
+    return gdrive_group_sync
+
+
+def _load_jira_doc_sync() -> DocSyncFuncType:
+    from ee.onyx.external_permissions.jira.doc_sync import jira_doc_sync
+
+    return jira_doc_sync
+
+
+def _load_jira_group_sync() -> GroupSyncFuncType:
+    from ee.onyx.external_permissions.jira.group_sync import jira_group_sync
+
+    return jira_group_sync
+
+
+def _load_censor_salesforce_chunks() -> CensoringFuncType:
+    from ee.onyx.external_permissions.salesforce.postprocessing import (
+        censor_salesforce_chunks,
+    )
+
+    return censor_salesforce_chunks
+
+
+def _load_sharepoint_doc_sync() -> DocSyncFuncType:
+    from ee.onyx.external_permissions.sharepoint.doc_sync import sharepoint_doc_sync
+
+    return sharepoint_doc_sync
+
+
+def _load_sharepoint_group_sync() -> GroupSyncFuncType:
+    from ee.onyx.external_permissions.sharepoint.group_sync import sharepoint_group_sync
+
+    return sharepoint_group_sync
+
+
+def _load_slack_doc_sync() -> DocSyncFuncType:
+    from ee.onyx.external_permissions.slack.doc_sync import slack_doc_sync
+
+    return slack_doc_sync
+
+
+def _load_teams_doc_sync() -> DocSyncFuncType:
+    from ee.onyx.external_permissions.teams.doc_sync import teams_doc_sync
+
+    return teams_doc_sync
+
+
 class DocSyncConfig(BaseModel):
     doc_sync_frequency: int
     doc_sync_func: DocSyncFuncType
@@ -96,60 +219,60 @@ _SOURCE_TO_SYNC_CONFIG: dict[DocumentSource, SyncConfig] = {
     DocumentSource.GOOGLE_DRIVE: SyncConfig(
         doc_sync_config=DocSyncConfig(
             doc_sync_frequency=DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=gdrive_doc_sync,
+            doc_sync_func=_lazy_doc_sync(_load_gdrive_doc_sync),
             initial_index_should_sync=True,
         ),
         group_sync_config=GroupSyncConfig(
             group_sync_frequency=GOOGLE_DRIVE_PERMISSION_GROUP_SYNC_FREQUENCY,
-            group_sync_func=gdrive_group_sync,
+            group_sync_func=_lazy_group_sync(_load_gdrive_group_sync),
             group_sync_is_cc_pair_agnostic=False,
         ),
     ),
     DocumentSource.CONFLUENCE: SyncConfig(
         doc_sync_config=DocSyncConfig(
             doc_sync_frequency=CONFLUENCE_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=confluence_doc_sync,
+            doc_sync_func=_lazy_doc_sync(_load_confluence_doc_sync),
             initial_index_should_sync=False,
         ),
         group_sync_config=GroupSyncConfig(
             group_sync_frequency=CONFLUENCE_PERMISSION_GROUP_SYNC_FREQUENCY,
-            group_sync_func=confluence_group_sync,
+            group_sync_func=_lazy_group_sync(_load_confluence_group_sync),
             group_sync_is_cc_pair_agnostic=True,
         ),
     ),
     DocumentSource.JIRA: SyncConfig(
         doc_sync_config=DocSyncConfig(
             doc_sync_frequency=JIRA_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=jira_doc_sync,
+            doc_sync_func=_lazy_doc_sync(_load_jira_doc_sync),
             initial_index_should_sync=True,
         ),
         group_sync_config=GroupSyncConfig(
             group_sync_frequency=JIRA_PERMISSION_GROUP_SYNC_FREQUENCY,
-            group_sync_func=jira_group_sync,
+            group_sync_func=_lazy_group_sync(_load_jira_group_sync),
             group_sync_is_cc_pair_agnostic=True,
         ),
     ),
     DocumentSource.CANVAS: SyncConfig(
         doc_sync_config=DocSyncConfig(
             doc_sync_frequency=CANVAS_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=canvas_doc_sync,
+            doc_sync_func=_lazy_doc_sync(_load_canvas_doc_sync),
             initial_index_should_sync=True,
         ),
         group_sync_config=GroupSyncConfig(
             group_sync_frequency=CANVAS_PERMISSION_GROUP_SYNC_FREQUENCY,
-            group_sync_func=canvas_group_sync,
+            group_sync_func=_lazy_group_sync(_load_canvas_group_sync),
             group_sync_is_cc_pair_agnostic=False,
         ),
     ),
     DocumentSource.BOX: SyncConfig(
         doc_sync_config=DocSyncConfig(
             doc_sync_frequency=BOX_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=box_doc_sync,
+            doc_sync_func=_lazy_doc_sync(_load_box_doc_sync),
             initial_index_should_sync=True,
         ),
         group_sync_config=GroupSyncConfig(
             group_sync_frequency=BOX_PERMISSION_GROUP_SYNC_FREQUENCY,
-            group_sync_func=box_group_sync,
+            group_sync_func=_lazy_group_sync(_load_box_group_sync),
             group_sync_is_cc_pair_agnostic=False,
         ),
     ),
@@ -158,32 +281,32 @@ _SOURCE_TO_SYNC_CONFIG: dict[DocumentSource, SyncConfig] = {
     DocumentSource.SLACK: SyncConfig(
         doc_sync_config=DocSyncConfig(
             doc_sync_frequency=SLACK_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=slack_doc_sync,
+            doc_sync_func=_lazy_doc_sync(_load_slack_doc_sync),
             initial_index_should_sync=True,
         ),
     ),
     DocumentSource.GMAIL: SyncConfig(
         doc_sync_config=DocSyncConfig(
             doc_sync_frequency=DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=gmail_doc_sync,
+            doc_sync_func=_lazy_doc_sync(_load_gmail_doc_sync),
             initial_index_should_sync=False,
         ),
     ),
     DocumentSource.GITHUB: SyncConfig(
         doc_sync_config=DocSyncConfig(
             doc_sync_frequency=GITHUB_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=github_doc_sync,
+            doc_sync_func=_lazy_doc_sync(_load_github_doc_sync),
             initial_index_should_sync=True,
         ),
         group_sync_config=GroupSyncConfig(
             group_sync_frequency=GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY,
-            group_sync_func=github_group_sync,
+            group_sync_func=_lazy_group_sync(_load_github_group_sync),
             group_sync_is_cc_pair_agnostic=False,
         ),
     ),
     DocumentSource.SALESFORCE: SyncConfig(
         censoring_config=CensoringConfig(
-            chunk_censoring_func=censor_salesforce_chunks,
+            chunk_censoring_func=_lazy_censoring(_load_censor_salesforce_chunks),
         ),
     ),
     DocumentSource.MOCK_CONNECTOR: SyncConfig(
@@ -198,19 +321,19 @@ _SOURCE_TO_SYNC_CONFIG: dict[DocumentSource, SyncConfig] = {
     DocumentSource.TEAMS: SyncConfig(
         doc_sync_config=DocSyncConfig(
             doc_sync_frequency=TEAMS_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=teams_doc_sync,
+            doc_sync_func=_lazy_doc_sync(_load_teams_doc_sync),
             initial_index_should_sync=True,
         ),
     ),
     DocumentSource.SHAREPOINT: SyncConfig(
         doc_sync_config=DocSyncConfig(
             doc_sync_frequency=SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=sharepoint_doc_sync,
+            doc_sync_func=_lazy_doc_sync(_load_sharepoint_doc_sync),
             initial_index_should_sync=True,
         ),
         group_sync_config=GroupSyncConfig(
             group_sync_frequency=SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY,
-            group_sync_func=sharepoint_group_sync,
+            group_sync_func=_lazy_group_sync(_load_sharepoint_group_sync),
             group_sync_is_cc_pair_agnostic=False,
         ),
     ),

backend/onyx/background/celery/apps/app_base.py

@@ -41,7 +41,6 @@ from onyx.configs.app_configs import (
 )
 from onyx.configs.constants import ONYX_CLOUD_CELERY_TASK_PREFIX, OnyxRedisLocks
 from onyx.db.engine.sql_engine import get_sqlalchemy_engine
-from onyx.document_index.opensearch.client import wait_for_opensearch_with_timeout
 from onyx.document_index.vespa.shared_utils.utils import wait_for_vespa_with_timeout
 from onyx.httpx.httpx_pool import HttpxPool
 from onyx.redis.redis_connector import RedisConnector
@@ -689,6 +688,11 @@ def wait_for_document_index_or_shutdown() -> None:
             raise WorkerShutdown(msg)
 
     if ENABLE_OPENSEARCH_INDEXING_FOR_ONYX:
+        # Imported here: opensearchpy costs ~18 MB and not every worker needs it.
+        from onyx.document_index.opensearch.client import (
+            wait_for_opensearch_with_timeout,
+        )
+
         if not wait_for_opensearch_with_timeout():
             msg = "[OpenSearch] Readiness probe did not succeed within the timeout. Exiting..."
             logger.error(msg)

backend/onyx/background/celery/apps/light.py

@@ -164,8 +164,6 @@ celery_app.autodiscover_tasks(
             "onyx.background.celery.tasks.doc_permission_syncing",
             "onyx.background.celery.tasks.docprocessing",
             "onyx.background.celery.tasks.opensearch_migration",
-            # Sandbox cleanup tasks (build feature)
-            "onyx.background.celery.tasks.build",
         ]
     )
 )

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

@@ -19,13 +19,6 @@ from onyx.server.features.build.db.sandbox import (
     user_has_stale_active_session,
 )
 from onyx.server.features.build.sandbox.factory import get_sandbox_manager
-from onyx.server.features.build.session.locks import get_session_creation_lock
-from onyx.server.features.build.session.sandbox_lifecycle import (
-    create_session_snapshot_keep_latest,
-    is_sandbox_idle,
-    list_snapshotable_session_workspaces,
-    sleep_sandbox,
-)
 
 # 100 minutes - snapshotting can take time
 TIMEOUT_SECONDS = 6000
@@ -54,6 +47,14 @@ def cleanup_idle_sandboxes_task(self: Task, *, tenant_id: str) -> None:  # noqa:
     fail-closed: snapshot failure on a reachable pod keeps the sandbox
     RUNNING for retry next sweep.
     """
+    from onyx.server.features.build.session.locks import get_session_creation_lock
+    from onyx.server.features.build.session.sandbox_lifecycle import (
+        create_session_snapshot_keep_latest,
+        is_sandbox_idle,
+        list_snapshotable_session_workspaces,
+        sleep_sandbox,
+    )
+
     task_logger.info(f"cleanup_idle_sandboxes_task starting for tenant {tenant_id}")
 
     redis_client = get_redis_client(tenant_id=tenant_id)

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

@@ -13,10 +13,6 @@ import logging
 from celery import Task, shared_task
 
 from onyx.background.celery.apps.app_base import task_logger
-from onyx.background.indexing.run_targeted_reindex import (
-    group_targets_by_cc_pair,
-    process_targets_for_cc_pair,
-)
 from onyx.configs.constants import OnyxCeleryTask
 from onyx.db.engine.sql_engine import get_session_with_current_tenant
 from onyx.db.enums import IndexingStatus
@@ -39,6 +35,11 @@ def run_targeted_reindex(
     """Body of the targeted-reindex task. Lifted out of the @shared_task
     decorator so tests can call it directly without going through
     celery's binding machinery."""
+    from onyx.background.indexing.run_targeted_reindex import (
+        group_targets_by_cc_pair,
+        process_targets_for_cc_pair,
+    )
+
     log = logging.LoggerAdapter(
         task_logger,
         extra={

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

@@ -109,7 +109,6 @@ from onyx.db.search_settings import (
     get_secondary_search_settings,
 )
 from onyx.db.swap_index import check_and_perform_index_swap
-from onyx.document_index.factory import get_all_document_indices
 from onyx.error_handling.exceptions import OnyxError
 from onyx.file_store.document_batch_storage import (
     DocumentBatchStorage,
@@ -117,19 +116,6 @@ from onyx.file_store.document_batch_storage import (
 )
 from onyx.file_store.staging import cleanup_staged_files_for_attempt
 from onyx.httpx.httpx_pool import HttpxPool
-from onyx.indexing.adapters.document_indexing_adapter import (
-    DocumentIndexingBatchAdapter,
-)
-from onyx.indexing.embedder import DefaultIndexingEmbedder
-from onyx.indexing.indexing_pipeline import run_indexing_pipeline
-from onyx.indexing.persistent_indexing import (
-    build_generic_connector_failure,
-    record_generic_failure,
-)
-from onyx.natural_language_processing.search_nlp_models import (
-    EmbeddingModel,
-    warm_up_bi_encoder,
-)
 from onyx.redis.redis_connector import RedisConnector
 from onyx.redis.redis_docprocessing import RedisDocprocessing
 from onyx.redis.redis_pool import (
@@ -878,6 +864,11 @@ def check_for_indexing(self: Task, *, tenant_id: str) -> int | None:
     w.r.t previous failed attempt, checkpointing, etc is handled in the docfetching task.
     """
 
+    from onyx.natural_language_processing.search_nlp_models import (
+        EmbeddingModel,
+        warm_up_bi_encoder,
+    )
+
     time_start = time.monotonic()
     task_logger.warning("check_for_indexing - Starting")
 
@@ -1552,6 +1543,11 @@ def _record_docprocessing_failure_persistent(
 
     Every step is wrapped so a follow-on error here does not re-raise out of
     the Celery task — we have already swallowed the original exception."""
+    from onyx.indexing.persistent_indexing import (
+        build_generic_connector_failure,
+        record_generic_failure,
+    )
+
     task_logger.info(
         "PERSISTENT_INDEXING enabled; recording docprocessing failure for "
         "attempt=%s batch=%s",
@@ -1710,6 +1706,14 @@ def _docprocessing_task(
     cross_batch_db_lock: RedisLock | None = None
 
     try:
+        # Inside the try so a failed first-use import still marks the attempt failed.
+        from onyx.document_index.factory import get_all_document_indices
+        from onyx.indexing.adapters.document_indexing_adapter import (
+            DocumentIndexingBatchAdapter,
+        )
+        from onyx.indexing.embedder import DefaultIndexingEmbedder
+        from onyx.indexing.indexing_pipeline import run_indexing_pipeline
+
         # FIX: Monitor memory before loading documents to track problematic batches
         emit_process_memory(
             os.getpid(),

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

@@ -11,8 +11,6 @@ from onyx.configs.app_configs import (
     SCHEDULED_EVAL_PROJECT,
 )
 from onyx.configs.constants import OnyxCeleryTask
-from onyx.evals.eval import run_eval
-from onyx.evals.models import EvalConfigurationOptions
 from onyx.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -31,6 +29,11 @@ def eval_run_task(
     configuration_dict: dict[str, Any],
 ) -> None:
     """Background task to run an evaluation with the given configuration"""
+    # Imported here: the eval stack pulls in braintrust, every tool and the chat
+    # pipeline (~90 MB), which the primary worker should not hold at boot.
+    from onyx.evals.eval import run_eval
+    from onyx.evals.models import EvalConfigurationOptions
+
     try:
         configuration = EvalConfigurationOptions.model_validate(configuration_dict)
         run_eval(configuration, remote_dataset_name=configuration.dataset_name)
@@ -90,6 +93,9 @@ def scheduled_eval_task(self: Task, **kwargs: Any) -> None:  # noqa: ARG001
         dataset_names,
     )
 
+    from onyx.evals.eval import run_eval
+    from onyx.evals.models import EvalConfigurationOptions
+
     pipeline_start = datetime.now(timezone.utc)
     results: list[dict[str, Any]] = []
 

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

@@ -50,9 +50,6 @@ from onyx.db.scheduled_task import (
     insert_run,
     mark_run_status,
 )
-from onyx.server.features.build.scheduled_tasks.executor import (
-    run_scheduled_task_logic,
-)
 from onyx.server.features.build.timeouts import (
     QUEUE_RESIDENCY_SECONDS,
     SCHEDULED_RUN_HARD_CAP_SECONDS,
@@ -233,6 +230,10 @@ def run_scheduled_task(self: Task, *, run_id: str, tenant_id: str) -> None:
     _ = self  # bound only for symmetry with other shared_task wrappers
     _ = tenant_id
     try:
+        from onyx.server.features.build.scheduled_tasks.executor import (
+            run_scheduled_task_logic,
+        )
+
         run_scheduled_task_logic(UUID(run_id))
     except Exception:
         task_logger.exception(

backend/onyx/db/engine/iam_auth.py

@@ -3,8 +3,6 @@ import os
 import ssl
 from typing import Any
 
-import boto3
-
 from onyx.configs.app_configs import (
     POSTGRES_HOST,
     POSTGRES_PORT,
@@ -20,6 +18,8 @@ def get_iam_auth_token(
     """
     Generate an IAM authentication token using boto3.
     """
+    import boto3
+
     client = boto3.client("rds", region_name=region)
     token = client.generate_db_auth_token(
         DBHostname=host, Port=int(port), DBUsername=user

backend/onyx/document_index/opensearch/client.py

@@ -7,7 +7,6 @@ from contextlib import AbstractContextManager, nullcontext
 from http import HTTPStatus
 from typing import Any, Generic, TypeVar
 
-import boto3
 from opensearchpy import (
     NotFoundError,
     OpenSearch,
@@ -294,6 +293,8 @@ class OpenSearchClient(AbstractContextManager):
             # SigV4 signing for an AWS managed domain whose FGAC master is an
             # IAM ARN. Credentials come from the default boto3 chain (env, IRSA,
             # instance/task role); the signer refreshes them per request.
+            import boto3
+
             credentials = boto3.Session().get_credentials()
             if credentials is None:
                 raise ValueError(

backend/onyx/evals/eval_cli.py

@@ -9,7 +9,6 @@ import logging
 import os
 from typing import Any
 
-import braintrust
 import requests
 
 from onyx.configs.app_configs import (
@@ -248,6 +247,8 @@ def main() -> None:
                 "--local-only cannot be used with --remote-dataset-name. Use --local-data-path with a local JSON file instead."
             )
         print(f"Loading data from remote dataset: {args.remote_dataset_name}")
+        import braintrust
+
         dataset = braintrust.init_dataset(
             project=args.braintrust_project, name=args.remote_dataset_name
         )

backend/onyx/file_store/file_store.py

@@ -5,9 +5,7 @@ from abc import ABC, abstractmethod
 from io import BytesIO
 from typing import IO, TYPE_CHECKING, Any, NotRequired, TypedDict, cast
 
-import boto3
 import puremagic
-from botocore.config import Config
 from botocore.exceptions import ClientError
 from sqlalchemy.orm import Session
 
@@ -222,6 +220,10 @@ class S3BackedFileStore(FileStore):
         """Initialize S3 client if not already done"""
         if self._s3_client is None:
             try:
+                # Imported here: boto3 costs ~16 MB and most workers never build an S3 client.
+                import boto3
+                from botocore.config import Config
+
                 client_kwargs: dict[str, Any] = {
                     "service_name": "s3",
                     "region_name": self._aws_region_name,

backend/tests/external_dependency_unit/db/test_targeted_reindex_task.py

@@ -44,8 +44,7 @@ from tests.external_dependency_unit.indexing_helpers import (
 )
 
 _PROCESSOR_PATH = (
-    "onyx.background.celery.tasks.docprocessing."
-    "targeted_reindex_task.process_targets_for_cc_pair"
+    "onyx.background.indexing.run_targeted_reindex.process_targets_for_cc_pair"
 )
 
 

backend/tests/external_dependency_unit/db/test_tenant_provisioning_rollback.py

@@ -79,12 +79,12 @@ class TestTenantProvisioningRollback:
             mock_redis.return_value.lock.return_value = mock_lock
 
             with patch(
-                "ee.onyx.background.celery.tasks.tenant_provisioning.tasks.setup_tenant"
+                "ee.onyx.server.tenants.provisioning.setup_tenant"
             ) as mock_setup:
                 mock_setup.side_effect = Exception("Simulated provisioning failure")
 
                 with patch(
-                    "ee.onyx.background.celery.tasks.tenant_provisioning.tasks.create_schema_if_not_exists",
+                    "ee.onyx.server.tenants.schema_management.create_schema_if_not_exists",
                     side_effect=track_schema_creation,
                 ):
                     # Run pre-provisioning - it should fail and trigger rollback

backend/tests/unit/onyx/document_index/opensearch/test_opensearch_auth.py

@@ -23,7 +23,7 @@ def test_iam_auth_uses_sigv4_signer() -> None:
     """
     with (
         patch("onyx.document_index.opensearch.client.OpenSearch") as mock_os,
-        patch("onyx.document_index.opensearch.client.boto3.Session") as mock_session,
+        patch("boto3.Session") as mock_session,
     ):
         mock_session.return_value.get_credentials.return_value = MagicMock()
         OpenSearchClient(
@@ -59,7 +59,7 @@ def test_iam_without_region_raises() -> None:
 
 
 def test_iam_without_credentials_raises() -> None:
-    with patch("onyx.document_index.opensearch.client.boto3.Session") as mock_session:
+    with patch("boto3.Session") as mock_session:
         mock_session.return_value.get_credentials.return_value = None
         with pytest.raises(ValueError, match="no AWS credentials"):
             OpenSearchClient(

tools/ods/internal/lazyimports/lazyimports.go

@@ -85,6 +85,18 @@ func DefaultLazyImportModules() map[string]LazyImportSettings {
 		"trafilatura":         NewLazyImportSettings(),
 		"pypdf":               NewLazyImportSettings(),
 		"unstructured_client": NewLazyImportSettings(),
+		"braintrust": NewLazyImportSettings(
+			"onyx/evals/providers/braintrust.py",
+			"onyx/tracing/braintrust_tracing_processor.py",
+		),
+		"exa_py": NewLazyImportSettings(
+			"onyx/tools/tool_implementations/web_search/clients/exa_client.py",
+		),
+		"playwright": NewLazyImportSettings(
+			"onyx/connectors/highspot/utils.py",
+			"onyx/connectors/web/connector.py",
+			"onyx/utils/playwright_fetch.py",
+		),
 	}
 }
 

tools/ods/internal/lazyimports/lazyimports_test.go

@@ -706,6 +706,9 @@ func TestDefaultLazyImportModules(t *testing.T) {
 		"trafilatura",
 		"pypdf",
 		"unstructured_client",
+		"braintrust",
+		"exa_py",
+		"playwright",
 	}
 
 	for _, mod := range expectedModules {
@@ -730,6 +733,71 @@ func TestDefaultLazyImportModules(t *testing.T) {
 			t.Errorf("Expected %s in litellm ignore files", ignore)
 		}
 	}
+
+	// Gateway modules are the only files allowed to import these at module level.
+	expectedGatewayIgnores := map[string][]string{
+		"braintrust": {
+			"onyx/evals/providers/braintrust.py",
+			"onyx/tracing/braintrust_tracing_processor.py",
+		},
+		"exa_py": {
+			"onyx/tools/tool_implementations/web_search/clients/exa_client.py",
+		},
+		"playwright": {
+			"onyx/connectors/highspot/utils.py",
+			"onyx/connectors/web/connector.py",
+			"onyx/utils/playwright_fetch.py",
+		},
+	}
+	for mod, ignores := range expectedGatewayIgnores {
+		for _, ignore := range ignores {
+			if _, ok := modules[mod].IgnoreFiles[ignore]; !ok {
+				t.Errorf("Expected %s in %s ignore files", ignore, mod)
+			}
+		}
+	}
+}
+
+func TestBraintrustAndExaViolations(t *testing.T) {
+	// Both packages pull in openai (and braintrust also litellm) at import time.
+	testContent := `
+import braintrust
+from braintrust import Eval
+from exa_py import Exa
+from exa_py.api import HighlightsContentsOptions
+import braintrust_langchain  # Different package, not flagged
+
+def allowed_function():
+    from braintrust import init_dataset
+    from exa_py import Exa
+    return init_dataset, Exa
+`
+
+	testPath := createTempPythonFile(t, testContent)
+	defer func() { _ = os.Remove(testPath) }()
+
+	patterns := createPatterns([]string{"braintrust", "exa_py"})
+	result, err := findEagerImports(testPath, patterns)
+	if err != nil {
+		t.Fatalf("findEagerImports failed: %v", err)
+	}
+
+	lineNumbers := extractLineNumbers(result.ViolationLines)
+	expectedLines := []int{2, 3, 4, 5}
+	if len(lineNumbers) != len(expectedLines) {
+		t.Errorf("Expected %d violations, got %d (lines %v)", len(expectedLines), len(lineNumbers), lineNumbers)
+	}
+	for _, expected := range expectedLines {
+		if !containsLineNum(lineNumbers, expected) {
+			t.Errorf("Expected line %d in violations", expected)
+		}
+	}
+
+	for _, mod := range []string{"braintrust", "exa_py"} {
+		if _, ok := result.ViolatedModules[mod]; !ok {
+			t.Errorf("Expected %s in violated modules", mod)
+		}
+	}
 }
 
 func TestIsValidPythonFile(t *testing.T) {