Automatically register DAGs that are used in a context manager #23592

ashb · 2022-05-09T16:05:15Z

This has caused a bit of confusion to committers (missing off the as dag) and is just more user friendly

This is on by default but can be disabled by passing auto_register=False do a DAG

ashb · 2022-05-09T16:06:35Z

This works for the current example sub-dags we have in the tests -- is there any other case I should test?

malthe · 2022-05-09T16:26:34Z

@ashb there are at least two things to consider here which I stumbled upon playing around with this:

There's nothing special about with except that tasks are automatically assigned to this dag (i.e., dag=dag). Tying registration of DAGs to with is perhaps a bit confusing in this sense – also, you can use with on a particular dag instance multiple times!
What about DAGs declared in functions?
```
def test():
    with DAG():
        ...
test()
```
I don't mind but it is a little magic perhaps, that this automatically declares the DAG.

I ended up instead with this pattern:

dag = DAG(...)

with dag:
    ...

– as a recommended pattern, simply because then you can't forget to give the DAG instance a top-level variable.

And then as I mention in https://lists.apache.org/thread/xndt3zklxbzo5nlmjz10dtfm2pp4t4wq, we could instead warn the user if they'd created a DAG which isn't exposed/registered.

(The gist of that linked reply is that there are two different ways to do that.)

kaxil · 2022-05-09T17:09:34Z

Tying registration of DAGs to with is perhaps a bit confusing in this sense – also, you can use with on a particular dag instance multiple times!

@malthe Is it the param name that is confusing? I do not understand I follow what you are saying. The use-case for this PR is that we don't need to needlessly do a with DAG(..) as dag as there is no need to have the as dag in it unless someone wants to actually use the dag variable itself

What about DAGs declared in functions?

Isn't it the same behaviour as the one with with DAG(..) as dag ?

malthe · 2022-05-09T19:30:54Z

@kaxil I understand the use-case – what I am somewhat hesitant to vote for is the conflation of the two orthogonal concerns which are both suggested for with here:

As a scope for tasks created which automatically get the dag parameter set;
As a mechanism to automatically register a DAG with the loading mechanism.

I think you could argue then that any DAG instance created during the loading of a module should be automatically registered – after all, why would a user create a DAG instance only to throw it away?

The problem with that is that we do have a test case that prevents this:

def test():
    dag = DAG(...)
    # Add tasks to dag and forget returning the DAG instance

test()

# or dag = test(), same result because `test` has no return value

The point is, should we require the use of with or should auto_register work regardless?

ashb · 2022-05-17T16:12:08Z

We could extend the auto-register -- but that makes it slightly more complex for SubDag (we'd have to "unregister" a dag when it gets turned in to a sub dag) but we could do it just as easily if we want to.

Edit: actually I guess that with subdag: is just as likely to be used as with dag

uranusjr · 2022-05-25T06:53:59Z

airflow/models/dag.py

@@ -3022,24 +3029,28 @@ class DagContext:

    """

-    _context_managed_dag: Optional[DAG] = None
-    _previous_context_managed_dags: List[DAG] = []
+    _context_managed_dags: Deque[DAG] = deque()


Any particular reason we need a deque here? From what I can tell a plain list is enough (with the ordering reversed, i.e. we append to / pop from the end and use -1 to access).

No particular reason no -- mostly thinking that we only ever need to access the head value, or push a new head value, and deque is more optimized for that case.

But it hardly matters as this is not on the critical path for performance.

uranusjr · 2022-08-01T06:20:27Z

tests/models/test_dagbag.py

@@ -937,8 +953,7 @@ def test_get_dag_with_dag_serialization(self):
    def test_collect_dags_from_db(self):
        """DAGs are collected from Database"""
        db.clear_db_dags()
-        example_dags_folder = airflow.example_dags.__path__[0]
-        dagbag = DagBag(example_dags_folder)
+        dagbag = DagBag(str(example_dags_folder))


Maybe we should just allow DagBag to take Path, but this doesn’t need to be done now.

Yeah, that was my thought -- it was orthogonal to this PR

newsfragments/23592.significant.rst

airflow/models/dag.py

potiuk

LGTM

ashb · 2022-08-05T08:56:21Z

I still need to update the docs/tutorial etc to reflect this.

ashb · 2022-08-12T09:50:48Z

Reminder to self: finish this before 2.4

airflow/example_dags/tutorial.py

uranusjr · 2022-09-06T19:41:46Z

airflow/models/dagbag.py

@@ -306,6 +314,8 @@ def _load_modules_from_file(self, filepath, safe_mode):
        if mod_name in sys.modules:
            del sys.modules[mod_name]

+        DagContext.current_autoregister_module_name = mod_name


This in-place operation feels a bit weird to me… Perhaps a context manager would be easier to understand and maintain? Something like

with DagContext.enable_autoregister(mod_name): parse(...)

This probably needs some refactoring to work though.

Yeah, I agree, This is essentially used as a global variable.

It's not very easy to refactor given it's use across three or so different functions, so a context manager would be tricky. (I.e. I can't see an easy refactor to make it work right now)

docs/apache-airflow/howto/dynamic-dag-generation.rst

newsfragments/23592.significant.rst

This has caused a bit of confusion to committers (missing off the `as dag`) and is just more user friendly This is on by default but can be disabled by passing `auto_register=False` to a DAG

ashb requested review from kaxil and XD-DENG as code owners May 9, 2022 16:05

ashb force-pushed the automatic-dag-registration branch from 4c2bdff to a62b484 Compare May 9, 2022 16:05

ashb requested review from malthe and potiuk May 9, 2022 16:09

ashb added this to the Airflow 2.4.0 milestone May 9, 2022

ashb force-pushed the automatic-dag-registration branch 2 times, most recently from c1a7796 to ac66a51 Compare May 24, 2022 09:21

uranusjr reviewed May 25, 2022

View reviewed changes

ashb force-pushed the automatic-dag-registration branch from ac66a51 to 301c20a Compare June 17, 2022 11:47

ashb force-pushed the automatic-dag-registration branch from 301c20a to bb5d2de Compare July 29, 2022 11:33

uranusjr reviewed Aug 1, 2022

View reviewed changes

uranusjr approved these changes Aug 1, 2022

View reviewed changes

ashb force-pushed the automatic-dag-registration branch from bb5d2de to ed6a98f Compare August 2, 2022 10:27

feluelle reviewed Aug 2, 2022

View reviewed changes

newsfragments/23592.significant.rst Show resolved Hide resolved

newsfragments/23592.significant.rst Show resolved Hide resolved

feluelle reviewed Aug 2, 2022

View reviewed changes

airflow/models/dag.py Show resolved Hide resolved

potiuk approved these changes Aug 4, 2022

View reviewed changes

ashb force-pushed the automatic-dag-registration branch from ed6a98f to 6cdd00e Compare August 15, 2022 11:20

uranusjr reviewed Aug 17, 2022

View reviewed changes

airflow/example_dags/tutorial.py Show resolved Hide resolved

potiuk approved these changes Aug 26, 2022

View reviewed changes

ashb force-pushed the automatic-dag-registration branch from 6cdd00e to 92fa142 Compare September 6, 2022 19:22

uranusjr reviewed Sep 6, 2022

View reviewed changes

jedcunningham reviewed Sep 7, 2022

View reviewed changes

ashb force-pushed the automatic-dag-registration branch from 291af76 to 2f793c5 Compare September 7, 2022 14:01

ashb added 6 commits September 7, 2022 17:07

Automatically register DAGs that are used in a context manager

b356357

This has caused a bit of confusion to committers (missing off the `as dag`) and is just more user friendly This is on by default but can be disabled by passing `auto_register=False` to a DAG

fixup! Automatically register DAGs that are used in a context manager

e99ff9c

fixup! Automatically register DAGs that are used in a context manager

f71c0dd

Cover dag decorator in newsfragment

8feae02

Update docs

39ffc24

fixup! Update docs

48f7281

ashb force-pushed the automatic-dag-registration branch from 72bbc61 to 48f7281 Compare September 7, 2022 16:07

jedcunningham approved these changes Sep 7, 2022

View reviewed changes

ashb merged commit 51c5d12 into apache:main Sep 7, 2022

ashb deleted the automatic-dag-registration branch September 7, 2022 17:54

ephraimbuddy added the type:new-feature Changelog: New Features label Sep 13, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Automatically register DAGs that are used in a context manager #23592

Automatically register DAGs that are used in a context manager #23592

ashb commented May 9, 2022 •

edited

ashb commented May 9, 2022

malthe commented May 9, 2022

kaxil commented May 9, 2022 •

edited

malthe commented May 9, 2022 •

edited

ashb commented May 17, 2022 •

edited

uranusjr May 25, 2022

ashb Jun 17, 2022

uranusjr Aug 1, 2022

ashb Aug 2, 2022

potiuk left a comment

ashb commented Aug 5, 2022

ashb commented Aug 12, 2022

uranusjr Sep 6, 2022 •

edited

ashb Sep 7, 2022

Automatically register DAGs that are used in a context manager #23592

Automatically register DAGs that are used in a context manager #23592

Conversation

ashb commented May 9, 2022 • edited

ashb commented May 9, 2022

malthe commented May 9, 2022

kaxil commented May 9, 2022 • edited

malthe commented May 9, 2022 • edited

ashb commented May 17, 2022 • edited

uranusjr May 25, 2022

Choose a reason for hiding this comment

ashb Jun 17, 2022

Choose a reason for hiding this comment

uranusjr Aug 1, 2022

Choose a reason for hiding this comment

ashb Aug 2, 2022

Choose a reason for hiding this comment

potiuk left a comment

Choose a reason for hiding this comment

ashb commented Aug 5, 2022

ashb commented Aug 12, 2022

uranusjr Sep 6, 2022 • edited

Choose a reason for hiding this comment

ashb Sep 7, 2022

Choose a reason for hiding this comment

ashb commented May 9, 2022 •

edited

kaxil commented May 9, 2022 •

edited

malthe commented May 9, 2022 •

edited

ashb commented May 17, 2022 •

edited

uranusjr Sep 6, 2022 •

edited