Add task instance mutation hook #8852
Conversation
|
@ashb could you take a look? Lyft leverages this features to do py2/py3 migration and found it pretty useful. I haven't followed actively on recent airflow development and not sure if there is anything similar exist in current codebase /feature. Let us know. |
| @@ -564,6 +565,27 @@ def with_all_tasks_removed(dag): | |||
| flaky_ti.refresh_from_db() | |||
| self.assertEqual(State.NONE, flaky_ti.state) | |||
|
|
|||
| @mock.patch('airflow.models.dagrun.task_instance_mutation_hook') | |||
There was a problem hiding this comment.
Can you add a test for TI where it is UP_FOR_RETRY
There was a problem hiding this comment.
@milton0825 Can you add a test for TI where it is UP_FOR_RETRY
There was a problem hiding this comment.
Testing that probably won't add too much value as the logic is not based on the state of TI.
There was a problem hiding this comment.
#8852 (comment) Based on our discussion in that comment if in future someone removes L468, no test would fail but it would break the behaviour that ti_mutation_hook won't work on ti that are up for retry.
|
|
||
| To define policy, add a ``airflow_local_settings`` module | ||
| to your PYTHONPATH that defines this ``policy`` function. It receives | ||
| a ``Task`` object and can alter it where needed. |
There was a problem hiding this comment.
Did you mean to remove these doc comments? This bit seems appropriate still.
There was a problem hiding this comment.
I moved that to the sections below.
…25/airflow into add-task-instance-mutation-hook
(cherry picked from commit bacb05d)
(cherry picked from commit bacb05d)
(cherry picked from commit bacb05d)
tl;dr: support task instance mutation before each execution.
Currently the only way to apply cluster-wide mutations to tasks is by defining
policyinairflow_local_setting. However,policyis applied totask(nottask_instanceand the descriptions and docs are wrong) right after the DAG is loaded. And you will not be able to change the task after that.Therefore a hook is added to support mutating the task instance before each execution (in case there is a retry). For example, you can re-route the task to execute in a different queue during retries.
Make sure to mark the boxes below before creating PR: [x]
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.
Read the Pull Request Guidelines for more information.