diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index 9e027e6c..5b17f6df 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -23,3 +23,5 @@ jobs: run: | pytest -vv --gherkin-terminal-reporter \ unity-test/unit + env: + PYTHONPATH: airflow/plugins diff --git a/airflow/dags/karpenter_test.py b/airflow/dags/karpenter_test.py index 7b03fa66..bd77ace9 100644 --- a/airflow/dags/karpenter_test.py +++ b/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, diff --git a/airflow/docker/custom_airflow/Dockerfile b/airflow/docker/custom_airflow/Dockerfile index c24bbf18..e11dc637 100644 --- a/airflow/docker/custom_airflow/Dockerfile +++ b/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 diff --git a/airflow/helm/values.tmpl.yaml b/airflow/helm/values.tmpl.yaml index a14d342e..0e85ce77 100644 --- a/airflow/helm/values.tmpl.yaml +++ b/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 diff --git a/airflow/plugins/__init__.py b/airflow/plugins/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/airflow/plugins/unity_sps_utils.py b/airflow/plugins/unity_sps_utils.py new file mode 100644 index 00000000..67b995f7 --- /dev/null +++ b/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