Automatically use BranchPythonOperator to handle if statements in TaskFlow DAG

[csm10495]

Say I have a DAG like this:

@dag(start_date=days_ago(1), schedule_interval=None, tags=["sample"])
def sample():
    @task()
    def get_random_value() -> int:
        return random.randint(5, 15)

    @task()
    def print_value(v):
        print(f"value: {v}")

    value = get_random_value()
    if value >= 7:
        print_value(value)

sample_dag = sample()

Right now it doesn't work since during render we get:

Broken DAG: [airflow/dags/sample.py] Traceback (most recent call last):
  File "site-packages/airflow/models/dag.py", line 2977, in factory
    f(**f_kwargs)
  File "airflow/dags/sample.py", line 44, in sample
    if value >= 7:
TypeError: '>=' not supported between instances of 'XComArg' and 'int'

It would be really cool if we could do branching based off of the results of tasks within TaskFlow DAGs. I guess internally it could use a PythonBranchOperator to figure out what should happen... or maybe some more fancy magic.

[potiuk]

Yeah., Would be great to make something like that works. But it would require some dark Python magic to happen I am afraid - due to the time/place the if is evaluated. Maybe posible though.

[csm10495]

@potiuk do we have a simple example of using BranchPythonOperator in taskflow (as it is today)?

I was playing around with some ast magic to see if i can find/replace if statements with branch operators (during @dag) but started hitting issues with BranchPythonOperator not being able to find tasks. I made it to here:

@dag(start_date=days_ago(1), schedule_interval=None, tags=["sample"])
def sample2():

    @task()
    def get_random_value() -> int:
        return random.randint(1, 3)

    @task()
    def print_value(v):
        print(f'value: {v}')

    @task()
    def print_value_again(v):
        print(f'value: {v}')

    @task()
    def body_1():
        print_value(value)

    @task()
    def or_else_1():
        pass

    def cond_1(arg):
        if True:
            return 'body_1'
        else:
            return 'or_else_1'

    value = get_random_value()
    print_value_again(value)
    BranchPythonOperator(task_id='auto_if_1', python_callable=cond_1, op_args=(value, ))

sample_dag2 = sample2()

[josh-fell]

Neither body_1 nor or_else_1 TaskFlow functions are being called (same for print_value but doesn't seem relevant for this particular example) so they won't appear in the DAG. Something like this based on the DAG above should do the trick:

@dag(start_date=datetime(2021, 1, 1), schedule_interval=None, tags=["sample"])
def sample2():

    @task()
    def get_random_value() -> int:
        return random.randint(1, 3)

    @task()
    def print_value(v):
        print(f'value: {v}')

    @task()
    def print_value_again(v):
        print(f'value: {v}')

    @task()
    def body_1():
        print_value(value)

    @task()
    def or_else_1():
        pass

    def cond_1(arg):
        if True:
            return 'body_1'
        else:
            return 'or_else_1'

    value = get_random_value()
    print_value_again(value)
    or_else = or_else_1()
    body1 = body_1()

    branch = BranchPythonOperator(task_id='auto_if_1', python_callable=cond_1, op_args=(value, ))
    branch >> [or_else, body1]


sample_dag2 = sample2()

[csm10495]

So i wrote some sample code that come close to being able to automatically turn if statements in taskflow dags into BranchPythonOperators:

"""
Sample of code that can rewrite a taskflow api dag to use BranchPythonOperator to replace if statements
"""

import ast

##
import functools
import inspect
import random

from textwrap import indent

from airflow.decorators import dag, task
from airflow.operators.python import BranchPythonOperator
from airflow.utils.dates import days_ago

try:
    from ast import unparse
except ImportError:
    from astunparse import unparse # pip install astunparse

def get_ast(thing, remove_decorators=True):
    source_lines = inspect.getsource(thing).splitlines()
    if remove_decorators:
        for idx, line in enumerate(source_lines):
            if line.startswith("def "):
                source_lines = source_lines[idx:]
                break

    source = "\n".join(source_lines)
    return ast.parse(source)


def _rewrite_this_function(function):
    class Rewrite(ast.NodeTransformer):
        def __init__(self):
            self.if_counter = 0

        def visit_If(self, node: ast.If):
            # todo: replace Node with a call to something else instead of If
            # print(ast.dump(node, indent=4))
            self.if_counter += 1
            module_text = f"""
@task()
def body_{self.if_counter}():
{indent(unparse(node.body).strip() or 'pass', '    ')}

@task()
def or_else_{self.if_counter}():
{indent(unparse(node.orelse).strip() or 'pass', '    ') }

def cond_{self.if_counter}(_locals):
    if eval("{unparse(node.test).strip()}", _locals):
        return 'body_{self.if_counter}'
    else:
        return 'or_else_{self.if_counter}'

bpo = BranchPythonOperator(task_id='auto_if_{self.if_counter}', python_callable=cond_{self.if_counter}, op_kwargs=dict(_locals=locals()))
bpo >> [body_{self.if_counter}(), or_else_{self.if_counter}()]

"""

            # print(module_text)

            new_node = ast.parse(module_text)
            new_node = ast.fix_missing_locations(new_node)
            new_node = ast.copy_location(new_node, new_node)

            # print(indent(unparse(new_node), "DEBUG: "))

            return new_node.body

    new_ast = Rewrite().visit(get_ast(function))
    new_ast_code = compile(new_ast, "<ast>", "exec")

    function_module = inspect.getmodule(function)
    exec(new_ast_code, function_module.__dict__)

    ret_func = getattr(function_module, function.__name__)

    print(unparse(new_ast))

    return ret_func


def rewrite_function(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        ret_func = _rewrite_this_function(func)
        return ret_func(*args, **kwargs)

    return wrapper


##


@dag(start_date=days_ago(1), schedule_interval=None, tags=["sample"])
@rewrite_function
def sample():
    @task()
    def get_random_value() -> int:
        return random.randint(1, 3)

    @task()
    def print_value(v):
        print(f"value (we passed!): {v}")

    @task()
    def the_value_is_2(v):
        print(f"Yay the value is 2! value: {v}")

    value = get_random_value()
    print_value(value)

    if value == 2:
        the_value_is_2(2)


sample_dag = sample()

It uses ast to find if-statements in the dag and then converts them into BranchPythonOperators. The weird thing is it converts the body and else of the if statement into their own Python Tasks. This has the unfortunate side effect of messing with the ability to call (or render) tasks from inside the body or else of the if statement.

So this dag:

@dag(start_date=days_ago(1), schedule_interval=None, tags=["sample"])
@rewrite_function
def sample():
    @task()
    def get_random_value() -> int:
        return random.randint(1, 3)

    @task()
    def print_value(v):
        print(f"value (we passed!): {v}")

    @task()
    def the_value_is_2(v):
        print(f"Yay the value is 2! value: {v}")

    value = get_random_value()
    print_value(value)

    if value == 2:
        the_value_is_2(2)


sample_dag = sample()

turns into:

@dag(start_date=days_ago(1), schedule_interval=None, tags=["sample"])
def sample():

    @task()
    def get_random_value() -> int:
        return random.randint(1, 3)

    @task()
    def print_value(v):
        print(f'value (we passed!): {v}')

    @task()
    def the_value_is_2(v):
        print(f'Yay the value is 2! value: {v}')
    value = get_random_value()
    print_value(value)

    @task()
    def body_1():
        the_value_is_2(2)

    @task()
    def or_else_1():
        pass

    def cond_1(_locals):
        if eval('(value == 2)', _locals):
            return 'body_1'
        else:
            return 'or_else_1'
    bpo = BranchPythonOperator(task_id='auto_if_1', python_callable=cond_1, op_kwargs=dict(_locals=locals()))
    (bpo >> [body_1(), or_else_1()])

It renders as expected (but without knowing about those inner task calls):

Seemingly pretty close, anyone have any ideas/pointers from here?

[vlavorini]

Hi! Thank you for the examples.

What If I want to branch over Task Groups? I tried with changing:

    # @task()
    # def or_else_1():
    #     pass
    @task_group()
    def or_else_1():
        @task()
        def mah():
            pass
        mah()

but I get the error:

Broken DAG: [/opt/airflow/dags/testtt.py] Traceback (most recent call last):
  File "/home/airflow/.local/lib/python3.7/site-packages/airflow/models/baseoperator.py", line 1474, in set_downstream
    self._set_relatives(task_or_task_list, upstream=False, edge_modifier=edge_modifier)
  File "/home/airflow/.local/lib/python3.7/site-packages/airflow/models/baseoperator.py", line 1412, in _set_relatives
    task_object.update_relative(self, not upstream)
AttributeError: 'NoneType' object has no attribute 'update_relative'