Add query job log correlation fields

Author: rishikeshdevsot

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

@@ -1,9 +1,20 @@
-from celery import Celery
+import logging
+
+from celery import Celery, signals
+from clp_py_utils.clp_logging import set_formatter_on_handlers
 
 from job_orchestration.executor.query import celeryconfig
 
 app = Celery("query")
 app.config_from_object(celeryconfig)
 
+
+@signals.after_setup_logger.connect
+@signals.after_setup_task_logger.connect
+def configure_logger_formatters(logger: logging.Logger | None = None, **_kwargs: object) -> None:
+    if logger is not None:
+        set_formatter_on_handlers(logger)
+
+
 if "__main__" == __name__:
     app.start()

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

@@ -14,7 +14,7 @@
     StorageType,
     WorkerConfig,
 )
-from clp_py_utils.clp_logging import set_logging_level
+from clp_py_utils.clp_logging import bind_log_context, set_logging_level
 from clp_py_utils.s3_utils import (
     generate_s3_url,
     get_credential_env_vars,
@@ -28,13 +28,20 @@
     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 ExtractIrJobConfig, ExtractJsonJobConfig
 from job_orchestration.scheduler.scheduler_data import QueryTaskStatus
 
 # Setup logging
 logger = get_task_logger(__name__)
 
 
+def _get_query_job_type(job_config: dict) -> str:
+    if "orig_file_id" in job_config:
+        return QueryJobType.EXTRACT_IR.name
+    return QueryJobType.EXTRACT_JSON.name
+
+
 def _make_clp_command_and_env_vars(
     clp_home: Path,
     worker_config: WorkerConfig,
@@ -201,7 +208,6 @@ def extract_stream_entry_point(
     logger.info(f"Started {task_name} task for job {job_id}")
 
     start_time = datetime.datetime.now()
-    task_status: QueryTaskStatus
     sql_adapter = SqlAdapter(Database.model_validate(clp_metadata_db_conn_params))
 
     # Load configuration
@@ -309,18 +315,26 @@ def extract_stream(
     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,
-        )
-    except SoftTimeLimitExceeded:
-        logger.exception(
-            f"Stream extraction task job_id={job_id} task_id={task_id} exceeded soft time limit."
-        )
-        raise
+    with bind_log_context(
+        job_id=job_id,
+        query_job_type=_get_query_job_type(job_config),
+        task_id=task_id,
+        archive_id=archive_id,
+        dataset=dataset,
+        celery_task_id=self.request.id,
+    ):
+        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(
+                f"Stream extraction task job_id={job_id} task_id={task_id} exceeded soft time limit."
+            )
+            raise

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

@@ -1,4 +1,5 @@
 import datetime
+import hashlib
 import os
 from pathlib import Path
 from typing import Any
@@ -13,7 +14,7 @@
     StorageType,
     WorkerConfig,
 )
-from clp_py_utils.clp_logging import set_logging_level
+from clp_py_utils.clp_logging import bind_log_context, set_logging_level
 from clp_py_utils.s3_utils import (
     generate_s3_url,
     get_credential_env_vars,
@@ -34,6 +35,13 @@
 logger = get_task_logger(__name__)
 
 
+def _get_query_log_context(query_string: str) -> dict[str, str]:
+    return {
+        "query": query_string,
+        "query_hash": hashlib.sha256(query_string.encode("utf-8")).hexdigest()[:16],
+    }
+
+
 def _make_core_clp_command_and_env_vars(
     clp_home: Path,
     worker_config: WorkerConfig,
@@ -243,47 +251,48 @@ def search_entry_point(
     clp_home = Path(os.getenv("CLP_HOME"))
     search_config = SearchJobConfig.model_validate(msgpack.unpackb(job_config_blob))
 
-    task_command, core_clp_env_vars = _make_command_and_env_vars(
-        clp_home=clp_home,
-        worker_config=worker_config,
-        archive_id=archive_id,
-        search_config=search_config,
-        results_cache_uri=results_cache_uri,
-        results_collection=job_id,
-        dataset=dataset,
-    )
-    if not task_command:
-        logger.error(f"Error creating {task_name} command")
-        return report_task_failure(
+    with bind_log_context(**_get_query_log_context(search_config.query_string)):
+        task_command, core_clp_env_vars = _make_command_and_env_vars(
+            clp_home=clp_home,
+            worker_config=worker_config,
+            archive_id=archive_id,
+            search_config=search_config,
+            results_cache_uri=results_cache_uri,
+            results_collection=job_id,
+            dataset=dataset,
+        )
+        if not task_command:
+            logger.error(f"Error creating {task_name} command")
+            return report_task_failure(
+                sql_adapter=sql_adapter,
+                task_id=task_id,
+                start_time=start_time,
+            )
+
+        task_results, _ = run_query_task(
             sql_adapter=sql_adapter,
+            logger=logger,
+            clp_logs_dir=clp_logs_dir,
+            task_command=task_command,
+            env_vars=core_clp_env_vars,
+            task_name=task_name,
+            job_id=job_id,
             task_id=task_id,
             start_time=start_time,
         )
 
-    task_results, _ = run_query_task(
-        sql_adapter=sql_adapter,
-        logger=logger,
-        clp_logs_dir=clp_logs_dir,
-        task_command=task_command,
-        env_vars=core_clp_env_vars,
-        task_name=task_name,
-        job_id=job_id,
-        task_id=task_id,
-        start_time=start_time,
-    )
-
-    storage_config = worker_config.stream_output.storage
-    if (
-        StorageType.S3 == storage_config.type
-        and search_config.write_to_file
-        and QueryTaskStatus.SUCCEEDED == task_results.status
-    ):
-        s3_config = storage_config.s3_config
-        src_file = Path(worker_config.stream_output.get_directory()) / job_id / archive_id
-        dest_path = f"{job_id}/{archive_id}"
-        upload_results_to_s3(task_results, s3_config, src_file, dest_path)
+        storage_config = worker_config.stream_output.storage
+        if (
+            StorageType.S3 == storage_config.type
+            and search_config.write_to_file
+            and QueryTaskStatus.SUCCEEDED == task_results.status
+        ):
+            s3_config = storage_config.s3_config
+            src_file = Path(worker_config.stream_output.get_directory()) / job_id / archive_id
+            dest_path = f"{job_id}/{archive_id}"
+            upload_results_to_s3(task_results, s3_config, src_file, dest_path)
 
-    return task_results.model_dump()
+        return task_results.model_dump()
 
 
 @app.task(bind=True)
@@ -297,16 +306,26 @@ 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,
-        )
-    except SoftTimeLimitExceeded:
-        logger.exception(f"Search task job_id={job_id} task_id={task_id} exceeded soft time limit.")
-        raise
+    with bind_log_context(
+        job_id=job_id,
+        query_job_type="SEARCH_OR_AGGREGATION",
+        task_id=task_id,
+        archive_id=archive_id,
+        dataset=dataset,
+        celery_task_id=self.request.id,
+    ):
+        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(
+                f"Search task job_id={job_id} task_id={task_id} exceeded soft time limit."
+            )
+            raise

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

@@ -9,6 +9,7 @@
 from typing import Any
 
 from clp_py_utils.clp_config import QUERY_TASKS_TABLE_NAME
+from clp_py_utils.clp_logging import bind_log_context
 from clp_py_utils.sql_adapter import SqlAdapter
 
 from job_orchestration.executor.utils import log_file_contents
@@ -53,6 +54,8 @@ def run_query_task(
 ) -> tuple[QueryTaskResult, str]:
     clo_log_path = get_task_log_file_path(clp_logs_dir, job_id, task_id)
     clo_log_file = open(clo_log_path, "w")
+    clo_log_file.write(f"job_id={job_id} task_id={task_id} task_name={task_name}\n")
+    clo_log_file.flush()
 
     task_status = QueryTaskStatus.RUNNING
     update_query_task_metadata(
@@ -69,34 +72,36 @@ def run_query_task(
         env=env_vars,
     )
 
-    def sigterm_handler(_signo, _stack_frame):
-        logger.debug("Entered sigterm handler")
-        if task_proc.poll() is None:
-            logger.debug(f"Trying to kill {task_name} process")
-            # Kill the process group in case the task process also forked
-            os.killpg(os.getpgid(task_proc.pid), signal.SIGTERM)
-            os.waitpid(task_proc.pid, 0)
-            logger.info(f"Cancelling {task_name} task.")
-        # Add 128 to follow convention for exit codes from signals
-        # https://tldp.org/LDP/abs/html/exitcodes.html#AEN23549
-        sys.exit(_signo + 128)
-
-    # Register the function to kill the child process at exit
-    signal.signal(signal.SIGTERM, sigterm_handler)
-
-    logger.info(f"Waiting for {task_name} to finish")
-    # `communicate` is equivalent to `wait` in this case, but avoids deadlocks when piping to
-    # stdout/stderr.
-    stdout_data, _ = task_proc.communicate()
-    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}"
-        )
-    else:
-        task_status = QueryTaskStatus.SUCCEEDED
-        logger.info(f"{task_name} task {task_id} completed for job {job_id}")
+    with bind_log_context(clp_subprocess_pid=task_proc.pid):
+
+        def sigterm_handler(_signo, _stack_frame):
+            logger.debug("Entered sigterm handler")
+            if task_proc.poll() is None:
+                logger.debug(f"Trying to kill {task_name} process")
+                # Kill the process group in case the task process also forked
+                os.killpg(os.getpgid(task_proc.pid), signal.SIGTERM)
+                os.waitpid(task_proc.pid, 0)
+                logger.info(f"Cancelling {task_name} task.")
+            # Add 128 to follow convention for exit codes from signals
+            # https://tldp.org/LDP/abs/html/exitcodes.html#AEN23549
+            sys.exit(_signo + 128)
+
+        # Register the function to kill the child process at exit
+        signal.signal(signal.SIGTERM, sigterm_handler)
+
+        logger.info(f"Waiting for {task_name} to finish")
+        # `communicate` is equivalent to `wait` in this case, but avoids deadlocks when piping to
+        # stdout/stderr.
+        stdout_data, _ = task_proc.communicate()
+        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}"
+            )
+        else:
+            task_status = QueryTaskStatus.SUCCEEDED
+            logger.info(f"{task_name} task {task_id} completed for job {job_id}")
 
     clo_log_file.close()
     if 0 != return_code: