From be4110ad92ce4061b8b09c75efe1ffc5c72cb445 Mon Sep 17 00:00:00 2001 From: James Ong Date: Tue, 18 Oct 2022 11:46:32 +0800 Subject: [PATCH 1/9] optimize task filtering --- airflow/models/taskinstance.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 54ce2a2d2e030..aaf97baccdf23 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2439,6 +2439,34 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau TaskInstance.task_id == first_task_id, ) + # group by dag_id and map_index while creating multiple groups for the other 2 params + if all(t.dag_id == dag_id and t.map_index == map_index for t in tis): + + # naively create 2 groups, 1 grouped by run_id and the other by task_id then use the smaller group + run_id_groups = defaultdict(list) + task_id_groups = defaultdict(list) + for t in tis: + run_id_groups[t.run_id].append(t.task_id) + task_id_groups[t.task_id].append(t.run_id) + filter_condition = [] + if len(run_id_groups) <= len(task_id_groups): + for run_id, task_ids in run_id_groups.items(): + filter_condition.append( + and_(TaskInstance.run_id == run_id, TaskInstance.task_id.in_(task_ids) + ) + ) + else: + for task_id, run_ids in task_id_groups.items(): + filter_condition.append( + and_(TaskInstance.task_id == task_id, TaskInstance.run_id.in_(run_ids) + ) + ) + return and_( + TaskInstance.dag_id == dag_id, + TaskInstance.map_index == map_index, + or_(*filter_condition) + ) + return tuple_in_condition( (TaskInstance.dag_id, TaskInstance.task_id, TaskInstance.run_id, TaskInstance.map_index), (ti.key.primary for ti in tis), From 51a0fb9e0daf11512bd2f8ef0ff577ffbcab8798 Mon Sep 17 00:00:00 2001 From: James Ong Date: Tue, 18 Oct 2022 17:21:11 +0800 Subject: [PATCH 2/9] generalize filtering --- airflow/models/taskinstance.py | 92 ++++++++++++++++++++-------------- 1 file changed, 55 insertions(+), 37 deletions(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index aaf97baccdf23..f412edfa5ac8e 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2417,60 +2417,78 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau first_task_id = first.task_id # Common path optimisations: when all TIs are for the same dag_id and run_id, or same dag_id # and task_id -- this can be over 150x faster for huge numbers of TIs (20k+) - if all(t.dag_id == dag_id and t.run_id == run_id and t.map_index == map_index for t in tis): + + # pre-compute the set of dag_id, run_id, map_indices and task_ids + dag_ids, run_ids, map_indices, task_ids = set(), set(), set(), set() + for t in tis: + dag_ids.add(t.dag_id) + run_ids.add(t.run_id) + map_indices.add(t.map_index) + task_ids.add(t.task_id) + + if dag_ids == {dag_id} and run_ids == {run_id} and map_indices == {map_index}: return and_( TaskInstance.dag_id == dag_id, TaskInstance.run_id == run_id, TaskInstance.map_index == map_index, - TaskInstance.task_id.in_(t.task_id for t in tis), + TaskInstance.task_id.in_(task_ids), ) - if all(t.dag_id == dag_id and t.task_id == first_task_id and t.map_index == map_index for t in tis): + if dag_ids == {dag_id} and task_ids == {first_task_id} and map_indices == {map_index}: return and_( TaskInstance.dag_id == dag_id, - TaskInstance.run_id.in_(t.run_id for t in tis), + TaskInstance.run_id.in_(run_ids), TaskInstance.map_index == map_index, TaskInstance.task_id == first_task_id, ) - if all(t.dag_id == dag_id and t.run_id == run_id and t.task_id == first_task_id for t in tis): + if dag_ids == {dag_id} and run_ids == {run_id} and task_ids == {task_id}: return and_( TaskInstance.dag_id == dag_id, TaskInstance.run_id == run_id, - TaskInstance.map_index.in_(t.map_index for t in tis), + TaskInstance.map_index.in_(map_indices), TaskInstance.task_id == first_task_id, ) - # group by dag_id and map_index while creating multiple groups for the other 2 params - if all(t.dag_id == dag_id and t.map_index == map_index for t in tis): - - # naively create 2 groups, 1 grouped by run_id and the other by task_id then use the smaller group - run_id_groups = defaultdict(list) - task_id_groups = defaultdict(list) - for t in tis: - run_id_groups[t.run_id].append(t.task_id) - task_id_groups[t.task_id].append(t.run_id) - filter_condition = [] - if len(run_id_groups) <= len(task_id_groups): - for run_id, task_ids in run_id_groups.items(): - filter_condition.append( - and_(TaskInstance.run_id == run_id, TaskInstance.task_id.in_(task_ids) - ) - ) - else: - for task_id, run_ids in task_id_groups.items(): - filter_condition.append( - and_(TaskInstance.task_id == task_id, TaskInstance.run_id.in_(run_ids) - ) - ) - return and_( - TaskInstance.dag_id == dag_id, - TaskInstance.map_index == map_index, - or_(*filter_condition) - ) + filter_condition = [] + # create 2 groups, 1 grouped by task_id the other by map_index and use the smaller group + task_id_groups, map_index_groups = defaultdict(lambda: defaultdict(list)), defaultdict(lambda: defaultdict(list)) + for t in tis: + task_id_groups[(t.dag_id, t.run_id)][t.task_id].append(t.map_index) + map_index_groups[(t.dag_id, t.run_id)][t.map_index].append(t.task_id) + + # this assumes that most dags have dag_id as the largest grouping, followed by run_id. even if its not, this is still + # a significant optimization over querying for every single tuple key + for cur_dag_id in dag_ids: + for cur_run_id in run_ids: + # we compare the group size between task_id and map_index and use the smaller group + dag_task_id_groups = task_id_groups[(cur_dag_id, cur_run_id)] + dag_map_index_groups = task_id_groups[(cur_dag_id, cur_run_id)] + + if len(dag_task_id_groups) <= len(dag_map_index_groups): + for cur_task_id, cur_map_indices in dag_task_id_groups.items(): + filter_condition.append( + and_( + TaskInstance.dag_id == cur_dag_id, + TaskInstance.run_id == cur_run_id, + TaskInstance.task_id == cur_task_id, + TaskInstance.map_index.in_(cur_map_indices) + ) + ) + else: + for cur_map_index, cur_task_ids in dag_map_index_groups.items(): + filter_condition.append( + and_( + TaskInstance.dag_id == cur_dag_id, + TaskInstance.run_id == cur_run_id, + TaskInstance.task_id.in_(cur_task_ids), + TaskInstance.map_index == cur_map_index, + ) + ) + + return and_( + TaskInstance.map_index == map_index, + or_(*filter_condition) + ) - return tuple_in_condition( - (TaskInstance.dag_id, TaskInstance.task_id, TaskInstance.run_id, TaskInstance.map_index), - (ti.key.primary for ti in tis), - ) @classmethod def ti_selector_condition(cls, vals: Collection[str | tuple[str, int]]) -> ColumnOperators: From 76b80e4774d4945df011169f17e3b708a14b6cb1 Mon Sep 17 00:00:00 2001 From: James Ong Date: Tue, 18 Oct 2022 17:21:55 +0800 Subject: [PATCH 3/9] fix bug on incorect vairable naming for task_id --- airflow/models/taskinstance.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index f412edfa5ac8e..cbaa2eabf4155 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2440,7 +2440,7 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau TaskInstance.map_index == map_index, TaskInstance.task_id == first_task_id, ) - if dag_ids == {dag_id} and run_ids == {run_id} and task_ids == {task_id}: + if dag_ids == {dag_id} and run_ids == {run_id} and task_ids == {first_task_id}: return and_( TaskInstance.dag_id == dag_id, TaskInstance.run_id == run_id, @@ -2484,10 +2484,7 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau ) ) - return and_( - TaskInstance.map_index == map_index, - or_(*filter_condition) - ) + return or_(*filter_condition) @classmethod From 0309336d8f91c58170770e00fa42005bded8b876 Mon Sep 17 00:00:00 2001 From: James Ong Date: Wed, 19 Oct 2022 13:35:23 +0800 Subject: [PATCH 4/9] add task mapping to existing test --- airflow/models/taskinstance.py | 4 ++-- tests/sensors/test_external_task_sensor.py | 21 ++++++++++++++++----- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index cbaa2eabf4155..5be0307ccdb80 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2415,8 +2415,6 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau run_id = first.run_id map_index = first.map_index first_task_id = first.task_id - # Common path optimisations: when all TIs are for the same dag_id and run_id, or same dag_id - # and task_id -- this can be over 150x faster for huge numbers of TIs (20k+) # pre-compute the set of dag_id, run_id, map_indices and task_ids dag_ids, run_ids, map_indices, task_ids = set(), set(), set(), set() @@ -2426,6 +2424,8 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau map_indices.add(t.map_index) task_ids.add(t.task_id) + # Common path optimisations: when all TIs are for the same dag_id and run_id, or same dag_id + # and task_id -- this can be over 150x faster for huge numbers of TIs (20k+) if dag_ids == {dag_id} and run_ids == {run_id} and map_indices == {map_index}: return and_( TaskInstance.dag_id == dag_id, diff --git a/tests/sensors/test_external_task_sensor.py b/tests/sensors/test_external_task_sensor.py index 80f538e8680ec..8aad10c67e845 100644 --- a/tests/sensors/test_external_task_sensor.py +++ b/tests/sensors/test_external_task_sensor.py @@ -24,6 +24,7 @@ import unittest import zipfile from datetime import time, timedelta +from airflow.decorators import task import pytest @@ -1116,6 +1117,10 @@ def dag_bag_head_tail(): dag_bag = DagBag(dag_folder=DEV_NULL, include_examples=False) with DAG("head_tail", start_date=DEFAULT_DATE, schedule="@daily") as dag: + @task + def dummy_task(x: int): + return x + head = ExternalTaskSensor( task_id="head", external_dag_id=dag.dag_id, @@ -1123,7 +1128,8 @@ def dag_bag_head_tail(): execution_delta=timedelta(days=1), mode="reschedule", ) - body = EmptyOperator(task_id="body") + + body = dummy_task.expand(x=[i for i in range(5)]) tail = ExternalTaskMarker( task_id="tail", external_dag_id=dag.dag_id, @@ -1153,15 +1159,20 @@ def test_clear_overlapping_external_task_marker(dag_bag_head_tail, session): ) session.add(dagrun) for task in dag.tasks: - ti = TaskInstance(task=task) - dagrun.task_instances.append(ti) + if task.task_id == "dummy_task": + for map_index in range(5): + ti = TaskInstance(task=task, map_index=map_index) + dagrun.task_instances.append(ti) + else: + ti = TaskInstance(task=task) + dagrun.task_instances.append(ti) ti.state = TaskInstanceState.SUCCESS session.flush() # The next two lines are doing the same thing. Clearing the first "head" with "Future" # selected is the same as not selecting "Future". They should take similar amount of # time too because dag.clear() uses visited_external_tis to keep track of visited ExternalTaskMarker. - assert dag.clear(start_date=DEFAULT_DATE, dag_bag=dag_bag_head_tail, session=session) == 30 + assert dag.clear(start_date=DEFAULT_DATE, dag_bag=dag_bag_head_tail, session=session) == 70 assert ( dag.clear( start_date=DEFAULT_DATE, @@ -1169,7 +1180,7 @@ def test_clear_overlapping_external_task_marker(dag_bag_head_tail, session): dag_bag=dag_bag_head_tail, session=session, ) - == 30 + == 70 ) From 9ac4b31b64f3953872638099635b9a85d0bd5037 Mon Sep 17 00:00:00 2001 From: James Ong Date: Tue, 1 Nov 2022 15:50:44 +0800 Subject: [PATCH 5/9] fix typo --- airflow/models/taskinstance.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 5be0307ccdb80..ad2f6190022a3 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2449,7 +2449,8 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau ) filter_condition = [] - # create 2 groups, 1 grouped by task_id the other by map_index and use the smaller group + # create 2 nested groups, both primarily grouped by dag_id and run_id, + # and in the nested group 1 grouped by task_id the other by map_index task_id_groups, map_index_groups = defaultdict(lambda: defaultdict(list)), defaultdict(lambda: defaultdict(list)) for t in tis: task_id_groups[(t.dag_id, t.run_id)][t.task_id].append(t.map_index) @@ -2461,7 +2462,7 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau for cur_run_id in run_ids: # we compare the group size between task_id and map_index and use the smaller group dag_task_id_groups = task_id_groups[(cur_dag_id, cur_run_id)] - dag_map_index_groups = task_id_groups[(cur_dag_id, cur_run_id)] + dag_map_index_groups = map_index_groups[(cur_dag_id, cur_run_id)] if len(dag_task_id_groups) <= len(dag_map_index_groups): for cur_task_id, cur_map_indices in dag_task_id_groups.items(): From b57f9b0f62a1098611a24c9258433054db4c7102 Mon Sep 17 00:00:00 2001 From: James Ong Date: Wed, 2 Nov 2022 21:55:13 +0800 Subject: [PATCH 6/9] remove test since recursively clearing mapped tasks does not work now --- tests/sensors/test_external_task_sensor.py | 21 +++++---------------- 1 file changed, 5 insertions(+), 16 deletions(-) diff --git a/tests/sensors/test_external_task_sensor.py b/tests/sensors/test_external_task_sensor.py index 8aad10c67e845..80f538e8680ec 100644 --- a/tests/sensors/test_external_task_sensor.py +++ b/tests/sensors/test_external_task_sensor.py @@ -24,7 +24,6 @@ import unittest import zipfile from datetime import time, timedelta -from airflow.decorators import task import pytest @@ -1117,10 +1116,6 @@ def dag_bag_head_tail(): dag_bag = DagBag(dag_folder=DEV_NULL, include_examples=False) with DAG("head_tail", start_date=DEFAULT_DATE, schedule="@daily") as dag: - @task - def dummy_task(x: int): - return x - head = ExternalTaskSensor( task_id="head", external_dag_id=dag.dag_id, @@ -1128,8 +1123,7 @@ def dummy_task(x: int): execution_delta=timedelta(days=1), mode="reschedule", ) - - body = dummy_task.expand(x=[i for i in range(5)]) + body = EmptyOperator(task_id="body") tail = ExternalTaskMarker( task_id="tail", external_dag_id=dag.dag_id, @@ -1159,20 +1153,15 @@ def test_clear_overlapping_external_task_marker(dag_bag_head_tail, session): ) session.add(dagrun) for task in dag.tasks: - if task.task_id == "dummy_task": - for map_index in range(5): - ti = TaskInstance(task=task, map_index=map_index) - dagrun.task_instances.append(ti) - else: - ti = TaskInstance(task=task) - dagrun.task_instances.append(ti) + ti = TaskInstance(task=task) + dagrun.task_instances.append(ti) ti.state = TaskInstanceState.SUCCESS session.flush() # The next two lines are doing the same thing. Clearing the first "head" with "Future" # selected is the same as not selecting "Future". They should take similar amount of # time too because dag.clear() uses visited_external_tis to keep track of visited ExternalTaskMarker. - assert dag.clear(start_date=DEFAULT_DATE, dag_bag=dag_bag_head_tail, session=session) == 70 + assert dag.clear(start_date=DEFAULT_DATE, dag_bag=dag_bag_head_tail, session=session) == 30 assert ( dag.clear( start_date=DEFAULT_DATE, @@ -1180,7 +1169,7 @@ def test_clear_overlapping_external_task_marker(dag_bag_head_tail, session): dag_bag=dag_bag_head_tail, session=session, ) - == 70 + == 30 ) From 891677aa979b289b111c1524cd72535f78aa53bf Mon Sep 17 00:00:00 2001 From: James Ong Date: Wed, 16 Nov 2022 14:58:04 -0500 Subject: [PATCH 7/9] fix typing --- airflow/models/taskinstance.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index ad2f6190022a3..164e9f2fd3382 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -11,7 +11,7 @@ # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF dict[dict[int, list[int]]] # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. @@ -2451,7 +2451,8 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau filter_condition = [] # create 2 nested groups, both primarily grouped by dag_id and run_id, # and in the nested group 1 grouped by task_id the other by map_index - task_id_groups, map_index_groups = defaultdict(lambda: defaultdict(list)), defaultdict(lambda: defaultdict(list)) + task_id_groups: dict[tuple, dict[Any, list[Any]]] = defaultdict(lambda: defaultdict(list)) + map_index_groups: dict[tuple, dict[Any, list[Any]]] = defaultdict(lambda: defaultdict(list)) for t in tis: task_id_groups[(t.dag_id, t.run_id)][t.task_id].append(t.map_index) map_index_groups[(t.dag_id, t.run_id)][t.map_index].append(t.task_id) From 59a355d484a6423cd1e0189792a882ca057d3ced Mon Sep 17 00:00:00 2001 From: James Ong Date: Wed, 16 Nov 2022 15:32:19 -0500 Subject: [PATCH 8/9] remove typo --- airflow/models/taskinstance.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 164e9f2fd3382..0adc510ebfe75 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -11,7 +11,7 @@ # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF dict[dict[int, list[int]]] +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. From 5bddbce5c507c7158764f369c6328bf15aa76c88 Mon Sep 17 00:00:00 2001 From: James Ong Date: Wed, 16 Nov 2022 17:10:09 -0500 Subject: [PATCH 9/9] run pre commit --- airflow/models/taskinstance.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 0adc510ebfe75..9ec2854be7ded 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2449,16 +2449,16 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau ) filter_condition = [] - # create 2 nested groups, both primarily grouped by dag_id and run_id, - # and in the nested group 1 grouped by task_id the other by map_index + # create 2 nested groups, both primarily grouped by dag_id and run_id, + # and in the nested group 1 grouped by task_id the other by map_index. task_id_groups: dict[tuple, dict[Any, list[Any]]] = defaultdict(lambda: defaultdict(list)) map_index_groups: dict[tuple, dict[Any, list[Any]]] = defaultdict(lambda: defaultdict(list)) for t in tis: task_id_groups[(t.dag_id, t.run_id)][t.task_id].append(t.map_index) map_index_groups[(t.dag_id, t.run_id)][t.map_index].append(t.task_id) - # this assumes that most dags have dag_id as the largest grouping, followed by run_id. even if its not, this is still - # a significant optimization over querying for every single tuple key + # this assumes that most dags have dag_id as the largest grouping, followed by run_id. even + # if its not, this is still a significant optimization over querying for every single tuple key for cur_dag_id in dag_ids: for cur_run_id in run_ids: # we compare the group size between task_id and map_index and use the smaller group @@ -2470,9 +2470,9 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau filter_condition.append( and_( TaskInstance.dag_id == cur_dag_id, - TaskInstance.run_id == cur_run_id, + TaskInstance.run_id == cur_run_id, TaskInstance.task_id == cur_task_id, - TaskInstance.map_index.in_(cur_map_indices) + TaskInstance.map_index.in_(cur_map_indices), ) ) else: @@ -2484,10 +2484,9 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> BooleanClau TaskInstance.task_id.in_(cur_task_ids), TaskInstance.map_index == cur_map_index, ) - ) - - return or_(*filter_condition) + ) + return or_(*filter_condition) @classmethod def ti_selector_condition(cls, vals: Collection[str | tuple[str, int]]) -> ColumnOperators: