How Airflow determines when to re-import a DAG file?

Viewed 66

I'm seeing some interesting behavior how Airflow treats an existing DAG file. For example, I have a DAG file like this:

# my_pipeline.py

from dag_gen import DagGenerator

g = globals()
g.update(DagGenerator.generate())

In the code above, DagGenerator.generate() returns a dictionary with two items:

{
    'my_dag': DAG(...),
    'my_op': DummyOperator(...),
}

Technically, this code should be equivalent to:

# my_pipeline_equivalent.py

my_dag = DAG(...)
my_op = DummyOperator(...)

For some reason, the file my_pipeline.py is not being picked up by Airflow while my_pipeline_equivalent.py can be picked up without any problem.

However, if I add the following code in my_pipeline.py, both dags (i.e. my_dag and my_dag2) can be picked up.

# my_pipeline.py

from dag_gen import DagGenerator

g = globals()
g.update(DagGenerator.generate())

# newly added
from airflow import DAG
from datetime import datetime
from airflow.operators.dummy_operator import DummyOperator

x = DAG(
    'my_dag2',
    schedule_interval="@daily",
    start_date=datetime(2021, 1, 1),
    catchup=False,
    default_args={
        "retries": 2, # If a task fails, it will retry 2 times.
    },
    tags=['example'],
)
b = DummyOperator(task_id='d3', dag=x)

What makes this stranger is if now I comment out the newly added part like below, my_dag can still be picked up while my_dag2 is gone.

# my_pipeline.py

from dag_gen import DagGenerator

g = globals()
g.update(DagGenerator.generate())

# newly added
#from airflow import DAG
#from datetime import datetime
#from airflow.operators.dummy_operator import DummyOperator

#x = DAG(
#    'my_dag2',
#    schedule_interval="@daily",
#    start_date=datetime(2021, 1, 1),
#    catchup=False,
#    default_args={
#        "retries": 2, # If a task fails, it will retry 2 times.
#    },
#    tags=['example'],
#)
#b = DummyOperator(task_id='d3', dag=x)

However, if I actually delete the commented code, my_dag is gone as well. In order to add my_dag back, I have to add back my_dag2 code without the comment (the commented my_dag2 code, which worked previously, doesn't work now).

Could anyone help me understand what's going here? If I remember correctly, Airflow has some logic to determine when/if to import/re-import a Python file. Does anyone know where that logic lives in code?

Thanks.

Some additional find-outs. Once I have both DAGs (my_dag and my_dag2) being picked up. If I change my_pipeline.py to only keep the DAG import, my_dag can still be picked up, even if I commented out the import, but I cannot remove that line. If I do remove it, my_dag would be gone again.

# my_pipeline.py

from dag_gen import DagGenerator

g = globals()
g.update(DagGenerator.generate())

# from airflow import DAG

My guess is Airflow must be reading the python file and looking for that import as string.

1 Answers

If you follow the steps that Airflow takes to process the files, you will get to this line of code:

return all(s in content for s in (b'dag', b'airflow'))

Which means that the Airflow file processor will ignore files that don't contain the two words dag and airflow.

If you want to process your module, and where you already have the word dag, you can just add a comment in the beginning contains the word airflow:

# airflow
from dag_gen import DagGenerator

g = globals()
g.update(DagGenerator.generate())

Or just renaming the class DagGenerator to AirflowDagGenerator to solve the problem in all the modules.

Related