k8s: added executor config

cern-sis · Apr 16, 2024 · 412ba71 · 412ba71
1 parent a126e4f
commit 412ba71
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 4 deletions.
diff --git a/dags/executor_config.py b/dags/executor_config.py
@@ -0,0 +1,14 @@
+from kubernetes import client as k8s
+
+kubernetes_executor_config = {
+    "pod_override": k8s.V1Pod(
+        spec=k8s.V1PodSpec(
+            containers=[
+                k8s.V1Container(
+                    name="base",
+                    image_pull_policy="Always",
+                )
+            ],
+        )
+    ),
+}
diff --git a/dags/open_access/gold_open_access_mechanisms.py b/dags/open_access/gold_open_access_mechanisms.py
@@ -4,6 +4,7 @@
 import pendulum
 from airflow.decorators import dag, task
 from airflow.providers.postgres.operators.postgres import PostgresOperator
+from executor_config import kubernetes_executor_config
 
 
 @dag(
@@ -12,7 +13,7 @@
     params={"year": 2023},
 )
 def oa_gold_open_access_mechanisms():
-    @task()
+    @task(executor_config=kubernetes_executor_config)
     def fetch_data_task(query, **kwargs):
         year = kwargs["params"].get("year")
         golden_access_base_query = (

diff --git a/dags/open_access/open_access.py b/dags/open_access/open_access.py
@@ -4,6 +4,7 @@
 import pendulum
 from airflow.decorators import dag, task
 from airflow.providers.postgres.operators.postgres import PostgresOperator
+from executor_config import kubernetes_executor_config
 
 
 @dag(
@@ -12,7 +13,7 @@
     params={"year": 2023},
 )
 def oa_dag():
-    @task()
+    @task(executor_config=kubernetes_executor_config)
     def fetch_data_task(query, **kwargs):
         year = kwargs["params"].get("year")
         base_query = (
@@ -26,7 +27,7 @@ def fetch_data_task(query, **kwargs):
         total = utils.get_total_results_count(data.text)
         return {type_of_query: total}
 
-    @task(multiple_outputs=True)
+    @task(multiple_outputs=True, executor_config=kubernetes_executor_config)
     def join(values, **kwargs):
         results = reduce(lambda a, b: {**a, **b}, values)
         results["years"] = kwargs["params"].get("year")

diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,4 @@
 -c https://raw.githubusercontent.com/apache/airflow/constraints-2.8.1/constraints-3.10.txt
-apache-airflow[celery, postgres, redis]==2.8.1
+apache-airflow[celery, postgres, redis, cncf.kubernetes]==2.8.1
 alembic==1.13.1
 airflow-provider-alembic==1.0.0