feat: migrate handler code from aind-airflow-dags (#20)

* build: require python >=3.9

* feat: add latest alert_handler

* test: update unit tests for alert_handler

* feat: add latest hpc_handler

* ci: add required paramiko version for appache-airflow-providers-ssh

* test: update unit tests for hpc_handler

* test: update unit tests for log_handler

* feat: add latest slurm_v2_handler with small modifications

* test: add unit tests for slurm_v2_handler (WIP)

* test: update unit tests for slurm_v2_handler methods and SlurmHook

* test: update unit tests for slurm_v2_handler.SubmitSlurmJobArray class

* fix: return null start/end times if no jobs found

* test: add unit tests for slurm_v2_handler.SlurmJobSensor class

* feat: remove aind-slurm-rest v1 handler and dependency

* feat: add latest task_handler

* test: update unit tests for task_handler

* refactor: linting

* ci: update Dockerfile to use alpine python base image

* ci: install optional loki and slurm dependencies in Dockerfile

Author: helen-m-lin

.github/workflows/test_and_lint.yml

@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        python-version: [ '3.8', '3.9', '3.10', '3.11' ]
+        python-version: [ '3.9', '3.10', '3.11' ]
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}

Dockerfile

@@ -1,8 +1,8 @@
-from python:3.10-slim
+FROM python:3.10-alpine
 
 WORKDIR /app
 ADD src ./src
 ADD pyproject.toml .
 ADD setup.py .
 
-RUN pip install . --no-cache-dir
+RUN pip install .[loki,slurm] --no-cache-dir

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 name = "aind-airflow-jobs"
 description = "Global classes for AIND Airflow service"
 license = {text = "MIT"}
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 authors = [
     {name = "Allen Institute for Neural Dynamics"}
 ]
@@ -42,15 +42,16 @@ all = [
 airflow = [
     "google-re2==1.0",
     "apache-airflow==2.8.1",
-    "apache-airflow-providers-ssh>=3.13.1,<=3.14.0"
+    "apache-airflow-providers-ssh>=3.13.1,<=3.14.0",
+    "paramiko>=2.9.0,<4.0.0"
 ]
 
 loki = [
     "python-logging-loki==0.3.1"
 ]
 
 slurm = [
-    "aind-slurm-rest==1.0.3"
+    "aind-slurm-rest-v2==0.0.3"
 ]
 
 [tool.setuptools.packages.find]
@@ -61,7 +62,7 @@ version = {attr = "aind_airflow_jobs.__version__"}
 
 [tool.black]
 line-length = 79
-target_version = ['py38']
+target_version = ['py39']
 exclude = '''
 
 (

src/aind_airflow_jobs/alert_handler.py

@@ -28,13 +28,18 @@ def action(self) -> str:
         }[self.value]
 
 
-def send_job_email(alert_type: AlertType, job: Dict[str, Any]) -> None:
+def send_job_email(
+    alert_type: AlertType,
+    job: Dict[str, Any],
+    task_id: Optional[str] = None,
+) -> None:
     """
-    Send an email given the alert type and job.
+    Send an email given the alert type, job, and optional task_id.
     Parameters
     ----------
     alert_type : AlertType
     job : Dict[str, Any]
+    task_id : Optional[str]
 
     Returns
     -------
@@ -49,10 +54,25 @@ def send_job_email(alert_type: AlertType, job: Dict[str, Any]) -> None:
         or alert_type.ALL.value in job.get("email_notification_types", [])
     ):
         to_email = job.get("user_email")
-        body = (
-            f"An airflow pipeline {reason} with the following conf\n\n"
-            f"Conf: {job}\n\n"
+        body = ""
+        if alert_type in [AlertType.FAIL, AlertType.RETRY]:
+            body += (
+                "Please check the AIND Data Transfer Service "
+                '<a href="http://aind-data-transfer-service/jobs">Job '
+                "Status Page</a> for more details or reach out to a "
+                "Scientific Computing engineer for assistance.<br/><br/>"
+            )
+
+        body += (
+            f"An airflow pipeline {reason} with the following "
+            f"configuration:<br/><br/>"
         )
+
+        if task_id:
+            body += f"Task: {task_id}<br/><br/>"
+
+        body += f"Configuration:<br/>{job}<br/><br/>"
+
         send_email(to=to_email, subject=subject, html_content=body)
 
 
@@ -135,7 +155,13 @@ def on_failure_or_retry_alert(
 
     """
     job = context["params"]
-    send_job_email(alert_type=alert_type, job=job)
+    task_id = context.get("task").task_id if context.get("task") else None
+
+    send_job_email(
+        alert_type=alert_type,
+        job=job,
+        task_id=task_id,
+    )
 
 
 def on_failure_or_retry_log_alert(
@@ -155,8 +181,14 @@ def on_failure_or_retry_log_alert(
     """
     job = context["params"]
     job_info = get_job_info_from_context(context)
+    task_id = context.get("task").task_id if context.get("task") else None
+
     send_log_message(job_info, log_level="ERROR")
-    send_job_email(alert_type=alert_type, job=job)
+    send_job_email(
+        alert_type=alert_type,
+        job=job,
+        task_id=task_id,
+    )
 
 
 def on_begin_or_end_alert(alert_type: AlertType, job: Dict[str, Any]) -> None:

src/aind_airflow_jobs/hpc_handler.py

@@ -3,6 +3,6 @@
 from airflow.providers.ssh.hooks.ssh import SSHHook
 
 
-def get_hpc_hook() -> SSHHook:
+def get_hpc_hook(ssh_conn_id: str = "hpc2/uri") -> SSHHook:
     """Returns a hook to send ssh commands to the hpc"""
-    return SSHHook(ssh_conn_id="hpc/uri")
+    return SSHHook(ssh_conn_id=ssh_conn_id)