Introduce TaskMixin #10930

turbaszek · 2020-09-14T11:32:30Z

Both BaseOperator and XComArgs implement bit shift operators
used to chain tasks in DAGs. By extracting this logic to new
mixin we reduce code duplication and make it easier to implement
it in the future.

Closes: #10926

^ Add meaningful description above

Read the Pull Request Guidelines for more information.
In case of fundamental code change, Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in UPDATING.md.

turbaszek · 2020-09-14T11:33:15Z

@casassg and @yuqian90 you may want to take a look 😉

turbaszek · 2020-09-14T12:43:21Z

Ok, this one is unexpected:

Too many ancestors (8/7) (too-many-ancestors)

What should we do:
A. Disable it locally in every affected operator
B. Increase limit of ancestors
C. Disable this check in pylintrc, as well as others like limit of arguments in class

Any opinions @mik-laj @kaxil @potiuk @ashb ?

kaxil · 2020-09-14T12:49:35Z

Ok, this one is unexpected:
Too many ancestors (8/7) (too-many-ancestors)
What should we do:
A. Disable it locally in every affected operator
B. Increase limit of ancestors
C. Disable this check in pylintrc, as well as others like limit of arguments in class

Any opinions @mik-laj @kaxil @potiuk @ashb ?

I would say B. increase the limit in pylintrc:

airflow/pylintrc

Lines 551 to 587 in eaa49b2

    
           [DESIGN] 
        
           # Maximum number of arguments for function / method. 
        
           # BasPH: choose 10 because this was 80% of the sorted list of number of arguments above 5 (Pylint default) 
        
           max-args=10 
        
           # Maximum number of attributes for a class (see R0902). 
        
           # BasPH: choose 15 because this was 80% of the sorted list of number of attributes above 7 (Pylint default) 
        
           max-attributes=15 
        
           # Maximum number of boolean expressions in an if statement. 
        
           max-bool-expr=5 
        
           # Maximum number of branch for function / method body. 
        
           # BasPH: choose 22 because this was 80% of the sorted list of number of attributes above 12 (Pylint default) 
        
           max-branches=22 
        
           # Maximum number of locals for function / method body. 
        
           # BasPH: choose 24 because this was 80% of the sorted list of number of locals above 15 (Pylint default) 
        
           max-locals=24 
        
           # Maximum number of parents for a class (see R0901). 
        
           max-parents=7 
        
           # Maximum number of public methods for a class (see R0904). 
        
           # BasPH: choose 27 because this was 50% of the sorted list of 30 number of public methods above 20 (Pylint default) 
        
           max-public-methods=27 
        
           # Maximum number of return / yield for function / method body. 
        
           max-returns=6 
        
           # Maximum number of statements in function / method body. 
        
           # BasPH: choose 69 because this was 80% of the sorted list of number of statements above 50 (Pylint default) 
        
           max-statements=69 
        
           # Minimum number of public methods for a class (see R0903). 
        
           min-public-methods=0

Currently, max-parents=7 we could increase it to 8

ashb · 2020-09-14T12:51:22Z

If vote b too.

yuqian90 · 2020-09-14T14:18:41Z

airflow/models/taskmixin.py

+    """
+
+    @abstractmethod
+    def set_upstream(self, other):


Do we want set_upstream to handle things such as List[XComArg] or List[TaskGroup]? If so, maybe we should add the typing to indicate other can be Union[TaskMixin, List[TaskMixin]].

I would love to add some type hints, but I'm not sure if using TaskMixin is that much informative. But I any other approach will create cyclic imports

airflow/models/taskmixin.py

casassg

Great! DRY!

kaxil · 2020-09-14T23:24:56Z

Static check is failing

yuqian90 · 2020-09-15T04:54:17Z

Thanks for the PR. This is a great improvement. I think there is room to "DRY" this further. For example, we have a lot of branching logic like this in these methods:

In BaseOperator._set_relatives(), we have this:

        if isinstance(task_or_task_list, XComArg):
            # otherwise we will start to iterate over xcomarg
            # because of the "list" check below
            # with current XComArg.__getitem__ implementation
            task_list = [task_or_task_list.operator]
        ...

When we add TaskGroup, we'll have a few more isinstance checks. We might be able to use some polymorphism to replace most of the if isinstance in these methods. One idea is to introduce a get_roots() and get_leaves() method within TaskMixin.

BaseOperator just returns itself as a list:

class BaseOperator(...):
    def get_roots(self) -> Generator["BaseOperator", None, None]:
        return [self]

    def get_leaves(self) -> Generator["BaseOperator", None, None]:
        return [self]

XComArgs returns the underlying operator:

class XComArg(...):
    def get_roots(self) -> Generator["BaseOperator", None, None]:
        return [self.operator]

    def get_leaves(self) -> Generator["BaseOperator", None, None]:
        return [self.operator]

TaskGroup has its own way of returning leaves or roots:

class TaskGroup(....):
    def get_roots(self) -> Generator["BaseOperator", None, None]:
       ...
    def get_leaves(self) -> Generator["BaseOperator", None, None]:
       ...

I think get_roots() and get_leaves() will be able to replace all the isinstance checks in set_upstream/set_downstream of XComArg, TaskGroup, BaseOperator.

What do you think?

turbaszek · 2020-09-15T08:52:08Z

I think get_roots() and get_leaves() will be able to replace all the isinstance checks in set_upstream/set_downstream of XComArg, TaskGroup, BaseOperator.

I must admit that I didn't get the get_roots and get_leaves idea fully. However, I implemented the operator property that should return BaseOperator and I think it serves the same purpose. WDYT @yuqian90 ?

yuqian90 · 2020-09-15T09:45:39Z

I think get_roots() and get_leaves() will be able to replace all the isinstance checks in set_upstream/set_downstream of XComArg, TaskGroup, BaseOperator.

I must admit that I didn't get the get_roots and get_leaves idea fully. However, I implemented the operator property that should return BaseOperator and I think it serves the same purpose. WDYT @yuqian90 ?

I see. For singular operators such as XComArg and BaseOperator, having a single operator property would suffice. However, when it comes to a group of operators with their own inter-dependencies, I don't think it's sufficient to return all the operators in the group. For example, when we do this, where operator1 is a singular operator, and group1 is a TaskGroup,

operator1 >> group1

We only want to set all the "roots" in group1 as downstream of operator1. So the operator1.set_downstream() somehow needs to have a way to get all the "roots" of group1. Here, "roots" means all the tasks within group1 that do not have any upstream dependencies within group1.

Same idea applies for "leaves" when we do this. Only the "leaves" in group1 needs to be pulled out and set upstream of operator1.

operator1 << group1

I think what you have at the moment (operator property) is good enough though for this PR's purpose. The "roots" vs "leaves" problem is very specific to TaskGroup. I'm fine with having such logic live within TaskGroup.

yuqian90 · 2020-09-16T00:37:01Z

airflow/models/baseoperator.py

@@ -1146,27 +1115,22 @@ def add_only_new(self, item_set: Set[str], item: str) -> None:
        else:
            item_set.add(item)

+    @property
+    def operator(self) -> "BaseOperator":


One request. Is it possible to make this property a List[BaseOperator] instead? If we want TaskGroup to implement TaskMixin, it'll be difficult to decide what operator should be if it's a singular operator.

I refactored it a bit so now its roots property - more in line with your initaill proposition

Both BaseOperator and XComArgs implement bit shift operators used to chain tasks in DAGs. By extracting this logic to new mixin we reduce code duplication and make it easier to implement it in future. Closes: apache#10926

yuqian90

Great. Thank you!

kaxil · 2020-09-16T10:22:44Z

airflow/example_dags/example_xcomargs.py

@@ -66,4 +66,4 @@ def print_value(value):
    xcom_args_a = print_value("first!")  # type: ignore
    xcom_args_b = print_value("second!")  # type: ignore

-    bash_op1 >> xcom_args_a >> xcom_args_b >> bash_op2
+    bash_op1 >> xcom_args_a >> xcom_args_b >> bash_op2  # type: ignore


oh ! is MyPy complaining here?

I don't get it and I would prefer to solve it in other PR as I believe this will require either mypy plugin or some type changes in @task decorator. The original error:

airflow/example_dags/example_xcomargs.py:66: error: Value of type variable "T" of "print_value" cannot be "str" xcom_args_a = print_value("first!") ^ airflow/example_dags/example_xcomargs.py:67: error: Value of type variable "T" of "print_value" cannot be "str" xcom_args_b = print_value("second!") ^ airflow/example_dags/example_xcomargs.py:69: error: Unsupported left operand type for >> ("BashOperator") bash_op1 >> xcom_args_a >> xcom_args_b >> bash_op2

ok, fixed? Use @task() instead of @task

hmm interesting, any idea here @casassg -- "@ task" should work.

This can be addressed in a separate PR though

Both @task and @task() works, however the type annotation of this decorator seems to be confusing mypy (althought it looks valid to me)

turbaszek · 2020-09-16T11:51:47Z

The single failing build seems to be a flaky something - @kaxil are we good with merging this PR?

kaxil · 2020-09-16T11:53:04Z

Yup 🚀

…ubDagOperator apache#10153 - Introduce TaskMixin (apache#10930) - Test updates

Both BaseOperator and XComArgs implement bit shift operators used to chain tasks in DAGs. By extracting this logic to new mixin we reduce code duplication and make it easier to implement it in the future. This change introduces also root property that allows us to reduce number of type checking. (cherry picked from 0779688)

turbaszek requested review from ashb and potiuk September 14, 2020 11:32

turbaszek added the area:providers label Sep 14, 2020

yuqian90 reviewed Sep 14, 2020

View reviewed changes

airflow/models/taskmixin.py Outdated Show resolved Hide resolved

casassg approved these changes Sep 14, 2020

View reviewed changes

turbaszek force-pushed the task-mixin branch from 2291ac2 to 792e1cf Compare September 15, 2020 08:53

yuqian90 reviewed Sep 16, 2020

View reviewed changes

turbaszek added 6 commits September 16, 2020 10:21

Introduce TaskMixin

a729e84

Both BaseOperator and XComArgs implement bit shift operators used to chain tasks in DAGs. By extracting this logic to new mixin we reduce code duplication and make it easier to implement it in future. Closes: apache#10926

Adjust pylintrc

0b091fd

fixup! Adjust pylintrc

7885ddb

Add operator property to TaskMixin

7fba960

Fix types

163a95c

Fix types

db38495

turbaszek force-pushed the task-mixin branch from ccd82cd to db38495 Compare September 16, 2020 08:22

turbaszek added 2 commits September 16, 2020 10:39

Rename operator to roots

47944c4

Typing in example

06a4080

yuqian90 approved these changes Sep 16, 2020

View reviewed changes

kaxil reviewed Sep 16, 2020

View reviewed changes

Typing in example

c5ee4a8

kaxil approved these changes Sep 16, 2020

View reviewed changes

turbaszek merged commit 0779688 into apache:master Sep 16, 2020

turbaszek deleted the task-mixin branch September 16, 2020 11:56

yuqian90 added a commit to yuqian90/airflow that referenced this pull request Sep 18, 2020

Use TaskMixin(apache#10930)

4c9464a

yuqian90 added a commit to yuqian90/airflow that referenced this pull request Sep 21, 2020

[AIP-34] TaskGroup: A UI task grouping concept as an alternative to S…

b6a80ec

…ubDagOperator apache#10153 - Introduce TaskMixin (apache#10930) - Test updates

yuqian90 mentioned this pull request Sep 24, 2020

[AIP-34] TaskGroup: A UI task grouping concept as an alternative to SubDagOperator #11119

Closed

turbaszek mentioned this pull request Oct 9, 2020

Airflow tutorial to use functional DAGs #11308

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Introduce TaskMixin #10930

Introduce TaskMixin #10930

turbaszek commented Sep 14, 2020

turbaszek commented Sep 14, 2020

turbaszek commented Sep 14, 2020

kaxil commented Sep 14, 2020 •

edited

Loading

ashb commented Sep 14, 2020

yuqian90 Sep 14, 2020

turbaszek Sep 14, 2020

casassg left a comment

kaxil commented Sep 14, 2020

yuqian90 commented Sep 15, 2020

turbaszek commented Sep 15, 2020

yuqian90 commented Sep 15, 2020

yuqian90 Sep 16, 2020

turbaszek Sep 16, 2020

yuqian90 left a comment

kaxil Sep 16, 2020

turbaszek Sep 16, 2020 •

edited

Loading

turbaszek Sep 16, 2020

kaxil Sep 16, 2020 •

edited

Loading

turbaszek Sep 16, 2020

turbaszek commented Sep 16, 2020

kaxil commented Sep 16, 2020

Introduce TaskMixin #10930

Introduce TaskMixin #10930

Conversation

turbaszek commented Sep 14, 2020

turbaszek commented Sep 14, 2020

turbaszek commented Sep 14, 2020

kaxil commented Sep 14, 2020 • edited Loading

ashb commented Sep 14, 2020

yuqian90 Sep 14, 2020

Choose a reason for hiding this comment

turbaszek Sep 14, 2020

Choose a reason for hiding this comment

casassg left a comment

Choose a reason for hiding this comment

kaxil commented Sep 14, 2020

yuqian90 commented Sep 15, 2020

turbaszek commented Sep 15, 2020

yuqian90 commented Sep 15, 2020

yuqian90 Sep 16, 2020

Choose a reason for hiding this comment

turbaszek Sep 16, 2020

Choose a reason for hiding this comment

yuqian90 left a comment

Choose a reason for hiding this comment

kaxil Sep 16, 2020

Choose a reason for hiding this comment

turbaszek Sep 16, 2020 • edited Loading

Choose a reason for hiding this comment

turbaszek Sep 16, 2020

Choose a reason for hiding this comment

kaxil Sep 16, 2020 • edited Loading

Choose a reason for hiding this comment

turbaszek Sep 16, 2020

Choose a reason for hiding this comment

turbaszek commented Sep 16, 2020

kaxil commented Sep 16, 2020

kaxil commented Sep 14, 2020 •

edited

Loading

turbaszek Sep 16, 2020 •

edited

Loading

kaxil Sep 16, 2020 •

edited

Loading