Add compression job log correlation fields

Author: rishikeshdevsot

components/job-orchestration/job_orchestration/executor/compress/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.compress import celeryconfig
 
 app = Celery("compress")
 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/compress/celery_compress.py

@@ -2,6 +2,7 @@
 from celery.app.task import Task
 from celery.exceptions import SoftTimeLimitExceeded
 from celery.utils.log import get_task_logger
+from clp_py_utils.clp_logging import bind_log_context
 
 from job_orchestration.executor.compress.celery import app
 from job_orchestration.executor.compress.compression_task import compression_entry_point
@@ -24,17 +25,22 @@ def compress(
     paths_to_compress_json: str,
     clp_metadata_db_connection_config,
 ):
-    try:
-        return compression_entry_point(
-            job_id,
-            task_id,
-            clp_io_config_json,
-            paths_to_compress_json,
-            clp_metadata_db_connection_config,
-            logger,
-        )
-    except SoftTimeLimitExceeded:
-        logger.exception(
-            f"Compression task job_id={job_id} task_id={task_id} exceeded soft time limit."
-        )
-        raise
+    with bind_log_context(
+        job_id=job_id,
+        task_id=task_id,
+        celery_task_id=self.request.id,
+    ):
+        try:
+            return compression_entry_point(
+                job_id,
+                task_id,
+                clp_io_config_json,
+                paths_to_compress_json,
+                clp_metadata_db_connection_config,
+                logger,
+            )
+        except SoftTimeLimitExceeded:
+            logger.exception(
+                f"Compression task job_id={job_id} task_id={task_id} exceeded soft time limit."
+            )
+            raise

components/job-orchestration/job_orchestration/executor/compress/compression_task.py

@@ -20,7 +20,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.clp_metadata_db_utils import (
     get_archives_table_name,
 )
@@ -432,8 +432,12 @@ def cleanup_temporary_files():
     # Open log files
     compression_log_path = logs_dir / f"{instance_id_str}-stderr.log"
     compression_log_file = open(compression_log_path, "w")
+    compression_log_file.write(f"job_id={job_id} task_id={task_id} task_name=compression\n")
+    compression_log_file.flush()
     conversion_log_path = logs_dir / f"{instance_id_str}-conversion-stderr.log"
     conversion_log_file = open(conversion_log_path, "w")
+    conversion_log_file.write(f"job_id={job_id} task_id={task_id} task_name=log-converter\n")
+    conversion_log_file.flush()
 
     conversion_return_code = 0
     if conversion_cmd is not None:
@@ -444,7 +448,8 @@ def cleanup_temporary_files():
             stderr=conversion_log_file,
             env=conversion_env,
         )
-        conversion_return_code = conversion_proc.wait()
+        with bind_log_context(clp_subprocess_pid=conversion_proc.pid):
+            conversion_return_code = conversion_proc.wait()
     conversion_log_file.close()
 
     try:
@@ -560,14 +565,15 @@ def cleanup_temporary_files():
             last_archive_stats = stats
 
         # Wait for compression to finish
-        return_code = proc.wait()
+        with bind_log_context(clp_subprocess_pid=proc.pid):
+            return_code = proc.wait()
 
-        if 0 != return_code:
-            logger.error(f"Failed to compress, return_code={return_code!s}")
-        else:
-            compression_successful = True
+            if 0 != return_code:
+                logger.error(f"Failed to compress, return_code={return_code!s}")
+            else:
+                compression_successful = True
 
-        logger.debug("Compressed.")
+            logger.debug("Compressed.")
     finally:
         cleanup_temporary_files()
         compression_log_file.close()

components/job-orchestration/job_orchestration/scheduler/compress/compression_scheduler.py

@@ -20,7 +20,7 @@
     OrchestrationType,
     StorageEngine,
 )
-from clp_py_utils.clp_logging import configure_logging, get_logger
+from clp_py_utils.clp_logging import bind_log_context, configure_logging, get_logger
 from clp_py_utils.clp_metadata_db_utils import (
     add_dataset,
     fetch_existing_datasets,
@@ -308,14 +308,15 @@ def search_and_schedule_new_tasks(
     # TODO: revisit why we need to commit here. To end long transactions?
     db_context.connection.commit()
     for job_row in jobs:
-        _schedule_job(
-            clp_config,
-            clp_metadata_db_connection_config,
-            task_manager,
-            db_context,
-            job_row,
-            existing_datasets,
-        )
+        with bind_log_context(job_id=job_row["id"]):
+            _schedule_job(
+                clp_config,
+                clp_metadata_db_connection_config,
+                task_manager,
+                db_context,
+                job_row,
+                existing_datasets,
+            )
 
 
 def _schedule_job(
@@ -474,54 +475,58 @@ def poll_running_jobs(
     logger.debug("Poll running jobs")
     jobs_to_delete = []
     for job_id, job in scheduled_jobs.items():
-        job_success = True
-        duration = 0.0
-        error_messages: list[str] = []
-        num_tasks_in_batch = 0
+        with bind_log_context(job_id=job_id):
+            job_success = True
+            duration = 0.0
+            error_messages: list[str] = []
+            num_tasks_in_batch = 0
 
-        try:
-            returned_results = job.result_handle.get_result()
-            if returned_results is None:
-                continue
+            try:
+                returned_results = job.result_handle.get_result()
+                if returned_results is None:
+                    continue
 
-            duration = (
-                datetime.datetime.now(datetime.timezone.utc) - job.start_time
-            ).total_seconds()
-            # Check for finished jobs
-            num_tasks_in_batch = len(returned_results)
-            for task_result in returned_results:
-                if task_result.status == CompressionTaskStatus.SUCCEEDED:
-                    logger.info(
-                        f"Compression task job-{job_id}-task-{task_result.task_id} completed in"
-                        f" {task_result.duration} second(s)."
-                    )
-                else:
-                    job_success = False
-                    error_messages.append(
-                        f"task {task_result.task_id}: {task_result.error_message}"
-                    )
-                    logger.error(
-                        f"Compression task job-{job_id}-task-{task_result.task_id} failed with"
-                        f" error: {task_result.error_message}."
-                    )
+                duration = (
+                    datetime.datetime.now(datetime.timezone.utc) - job.start_time
+                ).total_seconds()
+                # Check for finished jobs
+                num_tasks_in_batch = len(returned_results)
+                for task_result in returned_results:
+                    with bind_log_context(task_id=task_result.task_id):
+                        if task_result.status == CompressionTaskStatus.SUCCEEDED:
+                            logger.info(
+                                f"Compression task job-{job_id}-task-{task_result.task_id} completed in"
+                                f" {task_result.duration} second(s)."
+                            )
+                        else:
+                            job_success = False
+                            error_messages.append(
+                                f"task {task_result.task_id}: {task_result.error_message}"
+                            )
+                            logger.error(
+                                f"Compression task job-{job_id}-task-{task_result.task_id} failed with"
+                                f" error: {task_result.error_message}."
+                            )
 
-        except Exception:
-            logger.exception("Error while getting results for job %s", job_id)
-            job_success = False
+            except Exception:
+                logger.exception("Error while getting results for job %s", job_id)
+                job_success = False
 
-        if not job_success:
-            _handle_failed_compression_job(logs_directory, db_context, job_id, error_messages)
-            jobs_to_delete.append(job_id)
-            continue
+            if not job_success:
+                _handle_failed_compression_job(logs_directory, db_context, job_id, error_messages)
+                jobs_to_delete.append(job_id)
+                continue
 
-        job.num_tasks_completed += num_tasks_in_batch
+            job.num_tasks_completed += num_tasks_in_batch
 
-        if len(job.remaining_tasks) > 0:
-            _dispatch_next_task_batch(task_manager, db_context, job, max_concurrent_tasks_per_job)
-        else:
-            # All tasks completed successfully
-            _complete_compression_job(db_context, job_id, job.num_tasks_total, duration)
-            jobs_to_delete.append(job_id)
+            if len(job.remaining_tasks) > 0:
+                _dispatch_next_task_batch(
+                    task_manager, db_context, job, max_concurrent_tasks_per_job
+                )
+            else:
+                # All tasks completed successfully
+                _complete_compression_job(db_context, job_id, job.num_tasks_total, duration)
+                jobs_to_delete.append(job_id)
 
     for job_id in jobs_to_delete:
         del scheduled_jobs[job_id]