main changes

Author: rishikeshdevsot

components/job-orchestration/job_orchestration/executor/query/extract_stream_task.py

@@ -21,6 +21,7 @@
     s3_put,
 )
 from clp_py_utils.sql_adapter import SqlAdapter
+from structlog.contextvars import bound_contextvars
 
 from job_orchestration.executor.query.celery import app
 from job_orchestration.executor.query.utils import (
@@ -35,6 +36,25 @@
 logger = get_task_logger(__name__)
 
 
+def _get_extract_stream_task_log_context(
+    job_id: str,
+    task_id: int,
+    query_job_type: str,
+    archive_id: str,
+    dataset: str | None,
+) -> dict[str, Any]:
+    context: dict[str, Any] = {
+        "job_id": job_id,
+        "task_id": task_id,
+        "query_job_type": query_job_type,
+        "archive_id": archive_id,
+    }
+    if dataset is not None:
+        context["dataset"] = dataset
+
+    return context
+
+
 def _make_clp_command_and_env_vars(
     clp_home: Path,
     worker_config: WorkerConfig,
@@ -198,7 +218,7 @@ def extract_stream_entry_point(
     clp_logging_level = os.getenv("CLP_LOGGING_LEVEL")
     set_logging_level(logger, clp_logging_level)
 
-    logger.info(f"Started {task_name} task for job {job_id}")
+    logger.info(f"Started {task_name} task")
 
     start_time = datetime.datetime.now()
     task_status: QueryTaskStatus
@@ -303,24 +323,31 @@ def extract_stream(
     self: Task,
     job_id: str,
     task_id: int,
+    query_job_type: str,
     job_config: dict,
     archive_id: str,
     clp_metadata_db_conn_params: dict,
     results_cache_uri: str,
     dataset: str | None = None,
 ) -> dict[str, Any]:
-    try:
-        return extract_stream_entry_point(
-            job_id,
-            task_id,
-            job_config,
-            archive_id,
-            clp_metadata_db_conn_params,
-            results_cache_uri,
-            dataset,
+    with bound_contextvars(
+        **_get_extract_stream_task_log_context(
+            job_id, task_id, query_job_type, archive_id, dataset
         )
-    except SoftTimeLimitExceeded:
-        logger.exception(
-            f"Stream extraction task job_id={job_id} task_id={task_id} exceeded soft time limit."
-        )
-        raise
+    ):
+        try:
+            return extract_stream_entry_point(
+                job_id,
+                task_id,
+                job_config,
+                archive_id,
+                clp_metadata_db_conn_params,
+                results_cache_uri,
+                dataset,
+            )
+        except SoftTimeLimitExceeded:
+            logger.exception("Stream extraction task exceeded soft time limit.")
+            raise
+        except Exception:
+            logger.exception("Stream extraction task failed with an unexpected exception.")
+            raise

components/job-orchestration/job_orchestration/executor/query/fs_search_task.py

@@ -22,20 +22,45 @@
 )
 from clp_py_utils.sql_adapter import SqlAdapter
 from clp_py_utils.telemetry_config import is_telemetry_disabled_by_env
+from structlog.contextvars import bound_contextvars
 
 from job_orchestration.executor.query.celery import app
 from job_orchestration.executor.query.utils import (
+    get_query_hash,
     report_task_failure,
     run_query_task,
 )
 from job_orchestration.executor.utils import load_worker_config
+from job_orchestration.scheduler.constants import QueryJobType
 from job_orchestration.scheduler.job_config import SearchJobConfig
 from job_orchestration.scheduler.scheduler_data import QueryTaskResult, QueryTaskStatus
 
 # Setup logging
 logger = get_task_logger(__name__)
 
 
+def _get_search_task_log_context(
+    job_id: str,
+    task_id: int,
+    job_config_blob: bytes,
+    archive_id: str,
+    dataset: str | None,
+) -> dict[str, Any]:
+    context: dict[str, Any] = {
+        "job_id": job_id,
+        "task_id": task_id,
+        "query_job_type": QueryJobType.SEARCH_OR_AGGREGATION.to_str(),
+        "archive_id": archive_id,
+    }
+    if dataset is not None:
+        context["dataset"] = dataset
+
+    search_config = SearchJobConfig.model_validate(msgpack.unpackb(job_config_blob))
+    context["query"] = search_config.query_string
+    context["query_hash"] = get_query_hash(search_config.query_string)
+    return context
+
+
 def _make_core_clp_command_and_env_vars(
     clp_home: Path,
     worker_config: WorkerConfig,
@@ -238,7 +263,7 @@ def search_entry_point(
     clp_logging_level = os.getenv("CLP_LOGGING_LEVEL")
     set_logging_level(logger, clp_logging_level)
 
-    logger.info(f"Started {task_name} task for job {job_id}")
+    logger.info(f"Started {task_name} task")
 
     start_time = datetime.datetime.now()
     sql_adapter = SqlAdapter(Database.model_validate(clp_metadata_db_conn_params))
@@ -312,16 +337,24 @@ def search(
     results_cache_uri: str,
     dataset: str | None = None,
 ) -> dict[str, Any]:
-    try:
-        return search_entry_point(
-            job_id,
-            task_id,
-            job_config_blob,
-            archive_id,
-            clp_metadata_db_conn_params,
-            results_cache_uri,
-            dataset,
+    with bound_contextvars(
+        **_get_search_task_log_context(
+            job_id, task_id, job_config_blob, archive_id, dataset
         )
-    except SoftTimeLimitExceeded:
-        logger.exception(f"Search task job_id={job_id} task_id={task_id} exceeded soft time limit.")
-        raise
+    ):
+        try:
+            return search_entry_point(
+                job_id,
+                task_id,
+                job_config_blob,
+                archive_id,
+                clp_metadata_db_conn_params,
+                results_cache_uri,
+                dataset,
+            )
+        except SoftTimeLimitExceeded:
+            logger.exception("Search task exceeded soft time limit.")
+            raise
+        except Exception:
+            logger.exception("Search task failed with an unexpected exception.")
+            raise

components/job-orchestration/job_orchestration/executor/query/utils.py

@@ -1,4 +1,5 @@
 import datetime
+import hashlib
 import logging
 import os
 import signal
@@ -15,6 +16,10 @@
 from job_orchestration.scheduler.scheduler_data import QueryTaskResult, QueryTaskStatus
 
 
+def get_query_hash(query_string: str) -> str:
+    return hashlib.sha256(query_string.encode("utf-8")).hexdigest()
+
+
 def get_task_log_file_path(clp_logs_dir: Path, job_id: str, task_id: int) -> Path:
     worker_logs_dir = clp_logs_dir / job_id
     worker_logs_dir.mkdir(exist_ok=True, parents=True)
@@ -91,12 +96,10 @@ def sigterm_handler(_signo, _stack_frame):
     return_code = task_proc.returncode
     if 0 != return_code:
         task_status = QueryTaskStatus.FAILED
-        logger.error(
-            f"{task_name} task {task_id} failed for job {job_id} - return_code={return_code}"
-        )
+        logger.error(f"{task_name} task failed - return_code={return_code}")
     else:
         task_status = QueryTaskStatus.SUCCEEDED
-        logger.info(f"{task_name} task {task_id} completed for job {job_id}")
+        logger.info(f"{task_name} task completed")
 
     clo_log_file.close()
     if 0 != return_code: