Templating things in Airflow DAG
Follow into
Save into
Just another, TIL type entry. It is about Airflow. I have been working with Airflow for quite a while. I don't really like it to be frank, but I am stuck with it for now. With all that said, you can do everything that you want to do in Airflow. There is always some way to do everything. If you try really hard, you will be able to find even a way to capture Mewtwo. I am not saying this is something as important as that, but it took me a while to figure it out and I thought I should document it somewhere.
What I want to go over in this blog is ways in which we can get different variables into a dag which has KubernetesPodOperator.
Method 1 - From airflow config #
One easy place where we can pull a config from is the airflow.cfg file. This can be done with something like below:
1from airflow.configuration import conf
2
3with DAG(
4 dag_id="my-little-pony", schedule_interval=None, start_date=YESTERDAY
5) as dag:
6 head = KubernetesPodOperator(
7 image_pull_secrets=conf.get("kubernetes", "image_pull_secrets"),
8 image_pull_policy="Always",
9 name="build-head",
10 cmds=["python"],
11 arguments=["build-head"],
12 ...
13 )
If you see the example above, you can see how the image_pull_secrets value is fetched form the config. This works great if we just want to reuse things from the airflow config. I mainly use this for image_pull_secrets and namespace as they were same as the airflow webserver in my case.
Method 2 - Airflow Variables #
Next item is Variables. These are values that can be set from the Airflow UI. It is also possible to set them via env variables if we prefix the env var with AIRFLOW_VAR_. Relevant Airflow documentation.
In the UI, it is available under Admin>Variables. There you can create edit and delete and env var.
Once set, you can use an airflow variable like below:
1with DAG(
2 dag_id="my-little-pony", schedule_interval=None, start_date=YESTERDAY
3) as dag:
4 head = KubernetesPodOperator(
5 ...
6 cmds=["python"],
7 image=Variable.get("pony_builder_image_name"), # accessing variable
8 ...
9 )
Method 3 - Templating and macros #
Airflow actually uses Jinja templates for some stuff. It is kinda restricted, but it does provide some nice conveniences.
You can read more about it in the Airflow reference for Jinja Templating and macros.
I use this primary for values that are specific to the runtime of a airflow job. Things like run_id or ds(datetime stamp). You can also use this for two other things. You can use this to template out things from the config that gets passed in when you trigger a job from the UI or using the API. The config will be the json object and you can drill down any levels deep with just the .. You can also use this to template variables that I have mentioned earlier using var.
Here is a small example of where I use it.
1with DAG(
2 dag_id="my-little-pony", schedule_interval=None, start_date=YESTERDAY
3) as dag:
4 head = KubernetesPodOperator(
5 ...
6 cmds=["python"],
7 arguments=[
8 "{{ run_id }}", # use of run_id
9 "{{ dag_run.conf.name }}", # use of value from config
10 "{{ dag_run.conf.head.type }}", # use of nested config value
11 ],
12 ...
13 )
Method 4 - From env variables #
This is the main reason why I wanted to write this. This is mostly what I wanted to be able to do with my dag. The main reason why I got into this hunt in the first place is that I wanted to template the value of runAsUser etc from a kubernetes secret. I didn't ask for much man, but this lead to me to a relatively long hunt. Either my googling(ducking) skills are not so good or I was just not looking for the right thing. But anyways, I am here and I finally have solution.
This is what I wanted the code to end up looking like.
1with DAG(
2 dag_id="my-little-pony", schedule_interval=None, start_date=YESTERDAY
3) as dag:
4 head = KubernetesPodOperator(
5 ...
6 cmds=["python"],
7 arguments=[
8 "{{ run_id }}", # use of run_id
9 "{{ dag_run.conf.name }}", # use of value from config
10 "{{ dag_run.conf.head.type }}", # use of nested config value
11 ],
12 security_context={
13 "runAsUser": int(os.environ["PONY_RUN_AS_USER"]),
14 "runAsGroup": int(os.environ["PONY_RUN_AS_GROUP"]),
15 },
16 ...
17 )
So, if you see, things are relatively simple. I just load the env var from PONY_RUN_AS_USER and just use it. I have to convert it to int, but other than that, I just want to load it.
Quick first logical step would be to modify the airflow deployment yaml file to include these things when creating the image. This is mostly taken out of the yaml file that I was using. If you see, I am pulling the secrets from pony-secrets secret in my kube cluster and setting them as env variables.
1apiVersion: apps/v1
2kind: Deployment
3metadata:
4 name: airflow
5 namespace: pony-builder
6spec:
7 replicas: 1
8 selector:
9 matchLabels:
10 name: airflow
11 template:
12 metadata:
13 labels:
14 name: airflow
15 spec:
16 serviceAccountName: airflow
17 initContainers:
18 - name: "init"
19 image: pony-artifactory:latest
20 imagePullPolicy: Always
21 volumeMounts:
22 - name: airflow-configmap
23 mountPath: /root/airflow/airflow.cfg
24 subPath: airflow.cfg
25 env: &envvars
26 - name: PONY_RUN_AS_USER
27 valueFrom:
28 secretKeyRef:
29 name: pony-secrets
30 key: run_as_user
31 - name: PONY_RUN_AS_GROUP
32 valueFrom:
33 secretKeyRef:
34 name: pony-secrets
35 key: run_as_group
36 command:
37 - "bash"
38 args:
39 - "-cx"
40 - "/root/airflow-test-env-init.sh"
41 containers:
42 - name: webserver
43 image: pony-artifactory:latest
44 imagePullPolicy: Always
45 ports:
46 - name: webserver
47 containerPort: 80
48 args: ["webserver"]
49 envFrom:
50 - secretRef:
51 name: airflow
52 env: *envvars
53 volumeMounts:
54 - name: airflow-configmap
55 mountPath: /root/airflow/airflow.cfg
56 subPath: airflow.cfg
57 - name: scheduler
58 image: pony-artifactory:latest
59 imagePullPolicy: Always
60 envFrom:
61 - secretRef:
62 name: airflow
63 args: ["scheduler"]
64 env: *envvars
65 volumeMounts:
66 - name: airflow-configmap
67 mountPath: /root/airflow/airflow.cfg
68 subPath: airflow.cfg
69 volumes:
70 - name: airflow-configmap
71 configMap:
72 name: airflow-configmap
But, the issue with this is that these variables will not be available in the worker pods. I am not mentioning the final pods that gets launched, but when you launch the dag that I mentioned earlier, Airflow launches another pod which is the one that manages the actual pod. This intermediately pod does not have access to these env variables. Not sure what kind of a decision that was. I don't know if I am missing some kind of a flag or something that would enable this, but this was not available by default.
The job of this intermediate pod is to just evaluate the dag and then create a proper pod. So, if I have to evaluate the pod, I need access to these variables.
What happened here is that the webserver and the scheduler properly parsed the dags, but when this intermediate container that I mentioned tries to parse this dag config it fails saying that these env variables are not available.
What I ended up doing here is to add there to the kubernetes_secrets section in the configmap that gets mounted.
1[kubernetes_secrets]
2PONY_RUN_AS_USER = pony-secrets=run_as_user
3PONY_RUN_AS_GROUP = pony-secrets=run_as_group
Now with this, I can actually get it in the intermediate pod. And we can actually get the dag running. All good, finally. We can now finally make ponies.
Btw, this is how we make ponies. It was a well hidden secret, but now you know. It is all powered by Airflow and Kubernetes. The cuteness of gopher should have tipped you off that something was going on.
Bonus #
There is only other option. This is actually pretty well documented. If you don't really need the secret values to be used in the dag, but just want to have pass them as env variables, we can just use airflow.contrib.kubernetes.secret. Here is a sample code:
1from airflow.contrib.kubernetes import secret
2
3pony_secrets = [
4 secret.Secret(
5 deploy_type="env",
6 deploy_target="PONY_COLOR_MIX",
7 secret="pony-secrets",
8 key="color_mix",
9 ),
10 secret.Secret(
11 deploy_type="env",
12 deploy_target="PONY_TAIL_CURVE_RADIUS",
13 secret="pony-secrets",
14 key="tail_curve_radius",
15 ),
16]
17
18with DAG(
19 dag_id="my-little-pony", schedule_interval=None, start_date=YESTERDAY
20) as dag:
21 head = KubernetesPodOperator(
22 ...
23 cmds=["python"],
24 arguments=["build-head"],
25 secrets=pony_secrets, # use of secrets
26 ...
27 )
With that, I have most of what I had in my TIL bucket. Hope that was useful to someone.