thread safe approach to docprocessing logging

Author: evan-onyx

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

@@ -92,14 +92,14 @@
 from onyx.redis.redis_utils import is_fence
 from onyx.server.runtime.onyx_runtime import OnyxRuntime
 from onyx.utils.logger import setup_logger
-from onyx.utils.logger import TaskAttemptSingleton
 from onyx.utils.middleware import make_randomized_onyx_request_id
 from onyx.utils.telemetry import optional_telemetry
 from onyx.utils.telemetry import RecordType
 from shared_configs.configs import INDEXING_MODEL_SERVER_HOST
 from shared_configs.configs import INDEXING_MODEL_SERVER_PORT
 from shared_configs.configs import MULTI_TENANT
 from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
+from shared_configs.contextvars import INDEX_ATTEMPT_INFO_CONTEXTVAR
 
 logger = setup_logger()
 
@@ -1085,9 +1085,9 @@ def _docprocessing_task(
 ) -> None:
     start_time = time.monotonic()
 
-    # set the indexing attempt ID so that all log messages from this process
-    # will have it added as a prefix
-    TaskAttemptSingleton.set_cc_and_index_id(index_attempt_id, cc_pair_id)
+    # Cannot use the TaskSingleton approach here because the worker is multithreaded
+    INDEX_ATTEMPT_INFO_CONTEXTVAR.set((cc_pair_id, index_attempt_id))
+    task_logger
     if tenant_id:
         CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
 

backend/onyx/utils/logger.py

@@ -13,6 +13,7 @@
 from shared_configs.configs import SLACK_CHANNEL_ID
 from shared_configs.configs import TENANT_ID_PREFIX
 from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
+from shared_configs.contextvars import INDEX_ATTEMPT_INFO_CONTEXTVAR
 from shared_configs.contextvars import ONYX_REQUEST_ID_CONTEXTVAR
 
 
@@ -102,8 +103,14 @@ def process(
                     msg = f"[Doc Permissions Sync: {doc_permission_sync_ctx_dict['request_id']}] {msg}"
                 break
 
-            index_attempt_id = TaskAttemptSingleton.get_index_attempt_id()
-            cc_pair_id = TaskAttemptSingleton.get_connector_credential_pair_id()
+            index_attempt_info = INDEX_ATTEMPT_INFO_CONTEXTVAR.get()
+            if index_attempt_info:
+                cc_pair_id: int | None = index_attempt_info[0]
+                index_attempt_id: int | None = index_attempt_info[1]
+
+            else:
+                index_attempt_id = TaskAttemptSingleton.get_index_attempt_id()
+                cc_pair_id = TaskAttemptSingleton.get_connector_credential_pair_id()
 
             if index_attempt_id is not None:
                 msg = f"[Index Attempt: {index_attempt_id}] {msg}"

backend/shared_configs/contextvars.py

@@ -21,6 +21,11 @@
     "onyx_request_id", default=None
 )
 
+# Used to store cc pair id and index attempt id in multithreaded environments
+INDEX_ATTEMPT_INFO_CONTEXTVAR: contextvars.ContextVar[tuple[int, int] | None] = (
+    contextvars.ContextVar("index_attempt_info", default=None)
+)
+
 """Utils related to contextvars"""