Enable plugins folder

.github/workflows/unit_tests.yml

@@ -23,3 +23,5 @@ jobs:
       run: |
         pytest -vv --gherkin-terminal-reporter \
         unity-test/unit
+      env:
+        PYTHONPATH: airflow/plugins

airflow/dags/karpenter_test.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 
 from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
+from unity_sps_utils import get_affinity
 
 from airflow import DAG
 
@@ -15,65 +16,6 @@
 default_params = {"placeholder": 1}
 
 
-def get_affinity(
-    capacity_type: list[str], instance_family=list[str], instance_cpu=list[str], anti_affinity_label=str
-):
-
-    affinity = {
-        "nodeAffinity": {
-            "preferredDuringSchedulingIgnoredDuringExecution": [
-                {
-                    "weight": 1,
-                    "preference": {
-                        "matchExpressions": [
-                            {
-                                "key": "karpenter.sh/capacity-type",
-                                "operator": "In",
-                                "values": capacity_type,
-                            }
-                        ]
-                    },
-                }
-            ],
-            "requiredDuringSchedulingIgnoredDuringExecution": {
-                "nodeSelectorTerms": [
-                    {
-                        "matchExpressions": [
-                            {
-                                "key": "karpenter.k8s.aws/instance-family",
-                                "operator": "In",
-                                "values": instance_family,
-                            },
-                            {
-                                "key": "karpenter.k8s.aws/instance-cpu",
-                                "operator": "In",
-                                "values": instance_cpu,
-                            },
-                        ]
-                    }
-                ]
-            },
-        },
-        "podAntiAffinity": {
-            "requiredDuringSchedulingIgnoredDuringExecution": [
-                {
-                    "labelSelector": {
-                        "matchExpressions": [
-                            {
-                                "key": "app",
-                                "operator": "In",
-                                "values": [anti_affinity_label],
-                            },
-                        ]
-                    },
-                    "topologyKey": "kubernetes.io/hostname",
-                }
-            ]
-        },
-    }
-    return affinity
-
-
 dag = DAG(
     "karpenter_test",
     default_args=default_args,

airflow/docker/custom_airflow/Dockerfile

@@ -2,3 +2,5 @@ FROM apache/airflow:2.9.1-python3.11
 
 RUN pip install cwltool==3.1.20240112164112
 RUN pip install boto3==1.34.89
+
+COPY airflow/plugins /opt/airflow/plugins

airflow/helm/values.tmpl.yaml

@@ -278,6 +278,12 @@ env:
 
 # https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/security/api.html
 extraEnv: |
+  - name: AIRFLOW__CORE__DAGS_FOLDER
+    value: "/opt/airflow/dags"
+  - name: AIRFLOW__CORE__PLUGINS_FOLDER
+    value: "/opt/airflow/plugins"
+  - name: AIRFLOW__CORE__LAZY_LOAD_PLUGINS
+    value: "False"
   - name: AIRFLOW__API__AUTH_BACKENDS
     value: "airflow.api.auth.backend.basic_auth,airflow.api.auth.backend.session"
   - name: AIRFLOW__CORE__PARALLELISM

airflow/plugins/unity_sps_utils.py

@@ -0,0 +1,62 @@
+"""
+Module containing common utilities for the Unity Science Processing System.
+"""
+
+
+def get_affinity(
+    capacity_type: list[str], instance_family: list[str], instance_cpu: list[str], anti_affinity_label: str
+):
+
+    affinity = {
+        "nodeAffinity": {
+            "preferredDuringSchedulingIgnoredDuringExecution": [
+                {
+                    "weight": 1,
+                    "preference": {
+                        "matchExpressions": [
+                            {
+                                "key": "karpenter.sh/capacity-type",
+                                "operator": "In",
+                                "values": capacity_type,
+                            }
+                        ]
+                    },
+                }
+            ],
+            "requiredDuringSchedulingIgnoredDuringExecution": {
+                "nodeSelectorTerms": [
+                    {
+                        "matchExpressions": [
+                            {
+                                "key": "karpenter.k8s.aws/instance-family",
+                                "operator": "In",
+                                "values": instance_family,
+                            },
+                            {
+                                "key": "karpenter.k8s.aws/instance-cpu",
+                                "operator": "In",
+                                "values": instance_cpu,
+                            },
+                        ]
+                    }
+                ]
+            },
+        },
+        "podAntiAffinity": {
+            "requiredDuringSchedulingIgnoredDuringExecution": [
+                {
+                    "labelSelector": {
+                        "matchExpressions": [
+                            {
+                                "key": "app",
+                                "operator": "In",
+                                "values": [anti_affinity_label],
+                            },
+                        ]
+                    },
+                    "topologyKey": "kubernetes.io/hostname",
+                }
+            ]
+        },
+    }
+    return affinity