Conditional task

[kmehkeri]

I never liked the need of a separate branch task to conditionally execute another task and was thinking of a way to somehow decorate/enhance an arbitrary operator to skip it on condition.

Finally I came up with something that worked:

def RunIf(condition, operator):
    def wrapped_execute(self, context):
        if condition(context):
            self.original_execute(context)
        else:
            raise AirflowSkipException()
    operator.original_execute = operator.execute.__get__(operator, type(operator))
    operator.execute = wrapped_execute.__get__(operator, type(operator))
    return operator

dag = DAG('conditional_test', description='Conditional test', start_date=datetime(2017, 3, 20), catchup=False)
t1 = RunIf(lambda context: randint(0, 1) == 0, BashOperator(task_id='conditional1', dag=dag, bash_command="airflow version"))

Visually it looks okay to me and having run a simple test it looks to be working correctly...

But not being super advanced in Python I'm wondering if this is maybe too hacky, or prone to some errors or issues? Or not reasonable at all?

[potiuk]

I personally like it (and yes. I share the sentiment of not liking the separate branch operator/task).

Generally I think the concept is pretty sound and having a way that would conditionally skip execution of ANY operator (but otherwise not impacting the structure of the DAG) sounds rather appealing to me.

I think we could make it a little more sophisticated however, and nicer integrated into taskflow and with a little bit better syntactic sugar to make it look nicer. Simply, the example you show is just super-ugly and difficult to read.

Possibly we could even add it as a callable in the BaseOperator and any of our taskflow decorator that you could add so that it might look like:

Classic:

t1 = AnyOprator(task_id='task', run_if=lambda context: randint(0, 1) == 0)

TaskFlow:

@task.run_if(lambda context: randint(0, 1) == 0)
@task.python

Maybe naming could be a bit better and probably we could even pass context/map index into such "run_if` callable to make it way more powerful (and for example allow to skip execution of operators conditinally based on some of the context values and map index ).

I see no particular problem with that approach, I think it would be quite powerful to implement some "cross-cutting" skipping behaviour independent from the operator implementation.

There are probably some caveats:

• what to return as execute result if execute is not called. In this case (likely such run_if lambda should be able to return tuple - ("False". "return_value")

• decide when to run template mapping (ideally it should be done before such conditional check, which might be tricky if we want to make it part of another taskflow decorator

I wonder what others think (@dstandish @Taragolis @uranusjr @ephraimbuddy @josh-fell - calling the usual culprits who were usually implementing the decorators/syntactic sugar around the operator's use).

[dstandish]

I think it's a useful pattern and agree branch operators for this use case it's not great. That said, pretty sure one can already do this by using the pre_execute param in any operator (i.e. just pass your lambda to pre_execute). It already has access to the context and therefore can inspect upstream statuses, look at xcoms etc. WDYT?

[kmehkeri]

That's cool, I wasn't aware of pre_execute. So I could write something like:

def skip_at_random(context):
    if randint(0, 1) == 0:
        raise AirflowSkipException()

t2 = BashOperator(task_id='conditional2', pre_execute=skip_at_random, dag=dag, bash_command="airflow version")

And it is indeed much nicer than my initial version.

[potiuk]

Ah yeah - indeed. Would be nice to mention that pattern in the documentation. The best implementation of a feature is ... not to implement it at all :D

[kmehkeri]

I actually miss coherent documentation on branching. Currently it's a bit scattered around - some under DAGs/branching, but also an important mention in DAGs/trigger rules (which should be referenced from branching - because it's a common trap, looking at my colleagues and StackOverflow). Then there's separate operator how-to for BranchPythonOperator (and venv/external variants). And this what we're discussing here. I suppose we could come up with some other tricky scenarios.

How about a new section named "Branching and skipping tasks" under Authoring and Scheduling? If you think it's a good idea I could try contribute to this and submit PR for docs.