feat: follows logging standadrds in aind-log-utils

mekhlakapoor · mekhlakapoor · commit d604bca4eef8 · 2026-04-20T13:56:25.000-07:00
diff --git a/log_config_template.yaml b/log_config_template.yaml
@@ -0,0 +1,61 @@
+---
+version: 1
+disable_existing_loggers: false
+formatters:
+  default:
+    "()": aind_airflow_jobs.CustomJsonFormatter
+    format: "%(asctime)s %(levelname)s %(module)s %(lineno)s %(message)s"
+    reserved_attrs:
+      - args
+      - asctime
+      - created
+      - exc_info
+      - exc_text
+      - filename
+      - funcName
+      - levelname
+      - levelno
+      - lineno
+      - message
+      - module
+      - msecs
+      - msg
+      - name
+      - pathname
+      - process
+      - processName
+      - relativeCreated
+      - stack_info
+      - taskName
+      - thread
+      - threadName
+      - color_message
+    rename_fields:
+      "asctime": "timestamp"
+      "levelname": "level"
+    static_fields:
+      "environment": "local"
+      "software_version": ext://aind_airflow_jobs.__version__
+      "software_name": "aind-airflow-jobs"
+handlers:
+  console:
+    class: logging.StreamHandler
+    formatter: default
+    level: INFO
+    stream: ext://sys.stdout
+loggers:
+  "":
+    handlers:
+      - console
+    level: INFO
+    propagate: true
+  uvicorn.error:
+    handlers:
+      - console
+    level: INFO
+    propagate: false
+  uvicorn.access:
+    handlers:
+      - console
+    level: INFO
+    propagate: false
diff --git a/pyproject.toml b/pyproject.toml
@@ -17,6 +17,8 @@ readme = "README.md"
 dynamic = ["version"]
 
 dependencies = [
+    "log-json",
+    "PyYAML>=6.0,<7.0",
     "pydantic-settings>=2.2.0,<3.0",
     "pydantic>=2.2.0,<3.0"
 ]
diff --git a/src/aind_airflow_jobs/__init__.py b/src/aind_airflow_jobs/__init__.py
@@ -1,3 +1,44 @@
-"""Package to manage airflow jobs"""
+"""Package to manage airflow jobs."""
+
+import logging
+import logging.config
+import os
+from datetime import datetime, timezone
+from logging import LogRecord
+
+import log_json
+import yaml
 
 __version__ = "0.4.2"
+
+
+# We want to standardize the timestamp format to UTC and ISO-8601, which
+# requires a custom formatter and cannot be done through configuration only.
+class CustomJsonFormatter(log_json.JsonFormatter):
+    """Custom class to format log timestamps as ISO-8601 UTC."""
+
+    def formatTime(self, record: LogRecord, datefmt=None) -> str:
+        """
+        Format timestamp as ISO-8601 UTC.
+
+        Parameters
+        ----------
+        record : LogRecord
+        datefmt : str, optional
+            Unused parameter, kept for signature compatibility.
+
+        Returns
+        -------
+        str
+
+        """
+        dt = datetime.fromtimestamp(record.created, tz=timezone.utc)
+        return dt.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+
+
+if os.path.isfile(os.getenv("LOGGING_CONFIG_FILE", "log_config.yaml")):
+    config_path = os.getenv("LOGGING_CONFIG_FILE", "log_config.yaml")
+    with open(config_path, "rt", encoding="utf-8") as f:
+        config = yaml.safe_load(f.read())
+    logging.config.dictConfig(config)
+    logging.info("Found logging file at: %s", config_path)
diff --git a/src/aind_airflow_jobs/dag_tasks/base.py b/src/aind_airflow_jobs/dag_tasks/base.py
@@ -30,4 +30,10 @@ def run_task(self):
             )
 
         task_func()
-        logging.info(f"Task '{task_id}' completed successfully!")
+        logging.info(
+            f"Task '{task_id}' completed successfully!",
+            extra={
+                "process_name": task_id,
+                "pipeline_name": "airflow DAG",
+            },
+        )
diff --git a/src/aind_airflow_jobs/dag_tasks/check_connections.py b/src/aind_airflow_jobs/dag_tasks/check_connections.py
@@ -30,8 +30,20 @@ def check_param_store_connection(self):
         ams_uri = os.getenv("AMS_URI")
         co_uri = os.getenv("CO_URI")
 
-        logging.info(f"default_transfer_settings: {default_transfer_settings}")
-        logging.info(f"ams_uri: {ams_uri}")
+        logging.info(
+            f"default_transfer_settings: {default_transfer_settings}",
+            extra={
+                "process_name": "check_param_store_connection",
+                "pipeline_name": "airflow DAG",
+            },
+        )
+        logging.info(
+            f"ams_uri: {ams_uri}",
+            extra={
+                "process_name": "check_param_store_connection",
+                "pipeline_name": "airflow DAG",
+            },
+        )
 
         if not default_transfer_settings:
             raise AssertionError(
@@ -66,7 +78,13 @@ def check_slurm_connection(self):
         settings = SlurmClientSettings()
         slurm_api = settings.create_api_client()
         response = slurm_api.slurm_v0040_get_ping()
-        logging.info(f"SLURM ping response: {response}")
+        logging.info(
+            f"SLURM ping response: {response}",
+            extra={
+                "process_name": "check_slurm_connection",
+                "pipeline_name": "airflow DAG",
+            },
+        )
 
     def check_vast_connection(self):
         """Check that airflow can read files from VAST."""
@@ -89,7 +107,11 @@ def check_hpc_connection(self):
         ssh_command_output = self.airflow_task_settings.task_input_str or ""
         decoded = base64.b64decode(ssh_command_output).decode("utf-8")
         logging.info(
-            f"SSH Command Output: {ssh_command_output}. Decoded: {decoded}"
+            f"SSH Command Output: {ssh_command_output}. Decoded: {decoded}",
+            extra={
+                "process_name": "check_hpc_connection",
+                "pipeline_name": "airflow DAG",
+            },
         )
 
         if decoded.strip() != "Hello World":
diff --git a/src/aind_airflow_jobs/handlers/alert_handler.py b/src/aind_airflow_jobs/handlers/alert_handler.py
@@ -86,4 +86,13 @@ def send_log_message(
     level = level_name if isinstance(level_name, int) else logging.INFO
     logger = logging.getLogger(__name__)
     logger.setLevel(logging.INFO)
-    logger.log(level, sanitized_message)
+    logger.log(
+        level,
+        sanitized_message,
+        extra={
+            "acquisition_name": job_name,
+            "process_name": task_id,
+            "pipeline_name": "airflow DAG",
+            "run_id": run_id,
+        },
+    )
diff --git a/src/aind_airflow_jobs/handlers/slurm_v2_handler.py b/src/aind_airflow_jobs/handlers/slurm_v2_handler.py
@@ -347,7 +347,13 @@ def _requeue_failed_jobs(self, job_list: List[V0040JobInfo]) -> bool:
                     )
                     logging.warning(
                         f"Restarting {job.job_state} job: {job_id_to_retry}. "
-                        f"Restart count: {restart_count}"
+                        f"Restart count: {restart_count}",
+                        extra={
+                            "process_name": "submit_slurm_job_array",
+                            "pipeline_name": "airflow DAG",
+                            "job_id": job_id_to_retry,
+                            "job_state": job.job_state,
+                        },
                     )
                     sleep(1)
                     command = f"scontrol requeue {job_id_to_retry}"
@@ -463,11 +469,27 @@ def _monitor_job(
             last_backfill=V0040Uint64NoVal(),
             last_update=V0040Uint64NoVal(),
         )
-        logging.debug(f"Initialized placeholder: {job_response}")
+        logging.debug(
+            f"Initialized placeholder: {job_response}",
+            extra={
+                "process_name": "monitor_slurm_job",
+                "pipeline_name": "airflow DAG",
+                "acquisition_name": job_name,
+                "job_id": job_id,
+            },
+        )
         try:
             job_response = self.slurm.slurm_v0040_get_job(job_id=str(job_id))
         except NotFoundException:
-            logging.warning("Looking for job info in database...")
+            logging.warning(
+                "Looking for job info in database...",
+                extra={
+                    "process_name": "monitor_slurm_job",
+                    "pipeline_name": "airflow DAG",
+                    "acquisition_name": job_name,
+                    "job_id": job_id,
+                },
+            )
             slurm_db_api = SlurmdbApi(api_client=self.slurm.api_client)
             slurm_db_response = slurm_db_api.slurmdb_v0040_get_job(
                 job_id=str(job_id)
@@ -509,7 +531,15 @@ def _monitor_job(
                 "start_time": start_time,
             }
         )
-        logging.info(message)
+        logging.info(
+            message,
+            extra={
+                "process_name": "monitor_slurm_job",
+                "pipeline_name": "airflow DAG",
+                "acquisition_name": job_name,
+                "job_id": job_id,
+            },
+        )
         job_status = dict()
         is_finished, is_error = self._check_job_status(
             job_response, job_status
@@ -543,7 +573,15 @@ def _monitor_job(
                     "start_time": start_time,
                 }
             )
-            logging.info(message)
+            logging.info(
+                message,
+                extra={
+                    "process_name": "monitor_slurm_job",
+                    "pipeline_name": "airflow DAG",
+                    "acquisition_name": job_name,
+                    "job_id": job_id,
+                },
+            )
             is_finished, is_error = self._check_job_status(
                 job_response, job_status
             )
@@ -562,14 +600,30 @@ def _monitor_job(
                 remote_mnt_dir=self.remote_mnt_dir,
                 local_mnt_dir=self.local_mnt_dir,
             )
-            logging.exception(f"std_err:\n{std_err_msg}")
+            logging.exception(
+                f"std_err:\n{std_err_msg}",
+                extra={
+                    "process_name": "monitor_slurm_job",
+                    "pipeline_name": "airflow DAG",
+                    "acquisition_name": job_name,
+                    "job_id": job_id,
+                },
+            )
             raise Exception(
                 f"There were errors with the slurm job. "
                 f"Job: {message}. "
                 f"Errors: {errors}"
             )
         else:
-            logging.info("Job is Finished!")
+            logging.info(
+                "Job is Finished!",
+                extra={
+                    "process_name": "monitor_slurm_job",
+                    "pipeline_name": "airflow DAG",
+                    "acquisition_name": job_name,
+                    "job_id": job_id,
+                },
+            )
         return start_time, end_time
 
     def _std_err_filepath(self, job_id: Union[int, str]) -> str:
@@ -595,10 +649,34 @@ def run_job(self):
         submit_response = self._submit_job()
         job_id = submit_response.job_id
         job_name = self.job_properties.name
-        logging.info(f"Job Name: {job_name}")
-        logging.info(f"Job ID: {job_id}")
+        logging.info(
+            f"Job Name: {job_name}",
+            extra={
+                "process_name": "run_slurm_job",
+                "pipeline_name": "airflow DAG",
+                "acquisition_name": job_name,
+                "job_id": job_id,
+            },
+        )
+        logging.info(
+            f"Job ID: {job_id}",
+            extra={
+                "process_name": "run_slurm_job",
+                "pipeline_name": "airflow DAG",
+                "acquisition_name": job_name,
+                "job_id": job_id,
+            },
+        )
         std_err = self._std_err_filepath(job_id=job_id)
-        logging.info(f"Please check {std_err} for additional logs.")
+        logging.info(
+            f"Please check {std_err} for additional logs.",
+            extra={
+                "process_name": "run_slurm_job",
+                "pipeline_name": "airflow DAG",
+                "acquisition_name": job_name,
+                "job_id": job_id,
+            },
+        )
         self._monitor_job(submit_response=submit_response)
 
 
@@ -657,7 +735,13 @@ def _requeue_failed_jobs(self, job_list: List[V0040JobInfo]) -> bool:
                     )
                     logging.warning(
                         f"Restarting {job.job_state} job: {job_id_to_retry}. "
-                        f"Restart count: {restart_count}"
+                        f"Restart count: {restart_count}",
+                        extra={
+                            "process_name": "slurm_job_sensor",
+                            "pipeline_name": "airflow DAG",
+                            "job_id": job_id_to_retry,
+                            "job_state": job.job_state,
+                        },
                     )
                     sleep(1)
                     command = f"scontrol requeue {job_id_to_retry}"
@@ -751,15 +835,29 @@ def get_job_status(self) -> Tuple[bool, int, int]:
                     remote_mnt_dir=self.remote_mnt_dir,
                     local_mnt_dir=self.local_mnt_dir,
                 )
-                logging.exception(f"std_err:\n{std_err_msg}")
+                logging.exception(
+                    f"std_err:\n{std_err_msg}",
+                    extra={
+                        "process_name": "slurm_job_sensor",
+                        "pipeline_name": "airflow DAG",
+                        "job_id": self.job_id,
+                    },
+                )
                 raise Exception(
                     f"There were errors with the slurm job. "
                     f"Job: {self.job_id}. "
                     f"Errors: {errors}"
                 )
             return is_finished, start_time, end_time
         except NotFoundException:
-            logging.warning("Looking for job info in database...")
+            logging.warning(
+                "Looking for job info in database...",
+                extra={
+                    "process_name": "slurm_job_sensor",
+                    "pipeline_name": "airflow DAG",
+                    "job_id": job_id,
+                },
+            )
             slurm_db_api = SlurmdbApi(api_client=self.slurm.api_client)
             slurm_db_response = slurm_db_api.slurmdb_v0040_get_job(
                 job_id=str(job_id)
diff --git a/src/aind_airflow_jobs/handlers/task_handler.py b/src/aind_airflow_jobs/handlers/task_handler.py
diff --git a/tests/handlers/test_alert_handler.py b/tests/handlers/test_alert_handler.py
diff --git a/tests/handlers/test_slurm_v2_handler.py b/tests/handlers/test_slurm_v2_handler.py

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,8 @@ readme = "README.md"`
`17`	`17`	`dynamic = ["version"]`
`18`	`18`
`19`	`19`	`dependencies = [`
	`20`	`+ "log-json",`
	`21`	`+ "PyYAML>=6.0,<7.0",`
`20`	`22`	`"pydantic-settings>=2.2.0,<3.0",`
`21`	`23`	`"pydantic>=2.2.0,<3.0"`
`22`	`24`	`]`