You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2021/02/03 14:29:50 UTC
[GitHub] [airflow] yuqian90 opened a new pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
yuqian90 opened a new pull request #14048:
URL: https://github.com/apache/airflow/pull/14048
Clearing large number of tasks takes a long time. Most of the time is spent at this line in `clear_task_instances` (more than 95% time). This slowness sometimes causes the webserver to timeout because the `web_server_worker_timeout` is hit.
```python
# Clear all reschedules related to the ti to clear
session.query(TR).filter(
TR.dag_id == ti.dag_id,
TR.task_id == ti.task_id,
TR.execution_date == ti.execution_date,
TR.try_number == ti.try_number,
).delete()
```
This line is very slow because it's deleting `TaskReschedule` rows in a for loop one by one.
This PR simply changes this code to delete`TaskReschedule` in a single sql query with a bunch of `AND` conditions. It's effectively doing the same, but now it's much faster. Simple profiling shows that it's at least seven times faster when deleting thousands of `TaskReschedule`.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r570160429
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +170,23 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
Review comment:
Should we append `tis` and then use:
https://github.com/apache/airflow/blob/2a3960f359b1fcd5eed64fff472c4e463da74beb/airflow/models/taskinstance.py#L1987
to avoid code duplication?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] yuqian90 commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
yuqian90 commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r572006019
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +167,23 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
+
+ if tr_filter:
# Clear all reschedules related to the ti to clear
- session.query(TR).filter(
- TR.dag_id == ti.dag_id,
- TR.task_id == ti.task_id,
- TR.execution_date == ti.execution_date,
- TR.try_number == ti.try_number,
- ).delete()
+ delete_qry = TR.__table__.delete().where(
+ or_(
+ and_(
+ TR.dag_id == dag_id,
+ TR.task_id == task_id,
+ TR.execution_date == execution_date,
+ TR.try_number == try_number,
+ )
+ for dag_id, task_id, execution_date, try_number in tr_filter
Review comment:
Thanks. That's a great idea. I looked into the optimization in `filter_for_tis`. It works by extracting the common `dag_id` and `execution_date`, and then filter the `task_id` with an `IN` statement. This greatly reduces the size of the query, but it only does so if both `dag_id` and `execution_date` are equal for all tis.
For the case of slow `clear_task_instances`, the large number of tis is usually produced by the combination of a single dag_id over multiple `execution_date` for a large dag containing many task_ids. So it's slightly more complicated because there can be multiple `dag_id` and `execution_date` combinations (as well as `try_number`).
So I updated the logic to be more generic. It now uses a nested dictionary to produce the hierarchical filter. Some profiling showed great speed improvement (something like 40 to 50 times faster) compared to the first iteration. So the overral performance should now be 300 times faster than the original for loop deletion.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#issuecomment-772866641
The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest master at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
kaxil commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r571459217
##########
File path: airflow/models/taskinstance.py
##########
@@ -145,7 +145,11 @@ def clear_task_instances(
:param activate_dag_runs: flag to check for active dag run
:param dag: DAG object
"""
+ from sqlalchemy import and_, or_
Review comment:
this can go at the top of the file
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] yuqian90 commented on pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
yuqian90 commented on pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#issuecomment-775124153
Hi @kaxil this is the experiment I did. `test_create_large_dag_with_task_reschedule` creates a log of TaskInstance and `TaskReschedule` in the db. Then `test_clear_large_dag_with_task_reschedule` times how long it takes to call `clear_task_instances` on them.
```python
import pytest
@pytest.fixture(scope="module")
def big_dag():
with DAG(
'test_create_large_dag_with_task_reschedule',
start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE + datetime.timedelta(days=30),
) as dag:
for i in range(1000):
PythonSensor(task_id=f'task_{i}', python_callable=lambda: False, mode="reschedule")
yield dag
import copy
def test_create_large_dag_with_task_reschedule(big_dag):
tis = []
for i in range(10):
execution_date = DEFAULT_DATE + datetime.timedelta(days=i)
big_dag.create_dagrun(
execution_date=execution_date,
state=State.RUNNING,
run_type=DagRunType.SCHEDULED,
)
for task in big_dag.tasks:
ti = TI(task=copy.copy(task), execution_date=execution_date)
tis.append(ti)
import pendulum
tss = []
for ti in tis:
for i in range(5):
tss.append(TaskReschedule(
task=ti.task,
execution_date=ti.execution_date,
try_number=1,
start_date=pendulum.now(),
end_date=pendulum.now(),
reschedule_date=pendulum.now()))
def chunks(lst, n):
"""Yield successive n-sized chunks from lst."""
for i in range(0, len(lst), n):
yield lst[i:i + n]
with create_session() as session:
for chunk in chunks([{"task_id": tr.task_id, "dag_id": tr.dag_id, "execution_date": tr.execution_date,
"try_number": tr.try_number, "start_date": tr.start_date, "end_date": tr.end_date,
"duration": tr.duration, "reschedule_date": tr.reschedule_date} for tr in tss], 10000):
session.execute(TaskReschedule.__table__.insert(), chunk)
session.commit()
def test_clear_large_dag_with_task_reschedule(big_dag):
import time
with create_session() as session:
def count_task_reschedule():
return (
session.query(TaskReschedule)
.filter(
TaskReschedule.dag_id == big_dag.dag_id,
TaskReschedule.try_number == 1,
)
.count()
)
assert count_task_reschedule() == 1000 * 10 * 5
qry = session.query(TI).filter(TI.dag_id == big_dag.dag_id).all()
start = time.time()
clear_task_instances(qry, session, dag=big_dag)
end = time.time()
print(f"clear_task_instances took {end - start}")
assert count_task_reschedule() == 0
session.rollback()
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] yuqian90 commented on pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
yuqian90 commented on pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#issuecomment-773213408
Thanks @potiuk and @mik-laj for the review. I updated the PR according to your comments.
Also, I noticed that doing `TR.__table__.delete().where(...)` seems to be slightly faster than what I was doing with `session.query(...).filter(...).delete()` by a small margin. So I changed the PR to use it instead.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r571501423
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +167,23 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
+
+ if tr_filter:
# Clear all reschedules related to the ti to clear
- session.query(TR).filter(
- TR.dag_id == ti.dag_id,
- TR.task_id == ti.task_id,
- TR.execution_date == ti.execution_date,
- TR.try_number == ti.try_number,
- ).delete()
+ delete_qry = TR.__table__.delete().where(
+ or_(
+ and_(
+ TR.dag_id == dag_id,
+ TR.task_id == task_id,
+ TR.execution_date == execution_date,
+ TR.try_number == try_number,
+ )
+ for dag_id, task_id, execution_date, try_number in tr_filter
Review comment:
If these are all for the same dag (which it looks like they _should_ be given the API) then this query would be much quicker if `TR.dag_id = dag.dag_id` is pulled out of the loop so it doesn't appear so often.
(This was an optimization I made to `filter_for_tis`)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r572166901
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +167,23 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
+
+ if tr_filter:
# Clear all reschedules related to the ti to clear
- session.query(TR).filter(
- TR.dag_id == ti.dag_id,
- TR.task_id == ti.task_id,
- TR.execution_date == ti.execution_date,
- TR.try_number == ti.try_number,
- ).delete()
+ delete_qry = TR.__table__.delete().where(
+ or_(
+ and_(
+ TR.dag_id == dag_id,
+ TR.task_id == task_id,
+ TR.execution_date == execution_date,
+ TR.try_number == try_number,
+ )
+ for dag_id, task_id, execution_date, try_number in tr_filter
Review comment:
Yeah -- the performance difference is crazy for "large" queries in cases like this.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] yuqian90 commented on pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
yuqian90 commented on pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#issuecomment-775114357
Followed @ashb 's [suggestion](https://github.com/apache/airflow/pull/14048#discussion_r571501423) and speed it up even further. Please take another look.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] yuqian90 edited a comment on pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
yuqian90 edited a comment on pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#issuecomment-775124153
Hi @kaxil this is the experiment I did. `test_create_large_dag_with_task_reschedule` creates a lot of TaskInstance and `TaskReschedule` in the db. Then `test_clear_large_dag_with_task_reschedule` times how long it takes to call `clear_task_instances` on them.
```python
import pytest
@pytest.fixture(scope="module")
def big_dag():
with DAG(
'test_create_large_dag_with_task_reschedule',
start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE + datetime.timedelta(days=30),
) as dag:
for i in range(1000):
PythonSensor(task_id=f'task_{i}', python_callable=lambda: False, mode="reschedule")
yield dag
import copy
def test_create_large_dag_with_task_reschedule(big_dag):
tis = []
for i in range(10):
execution_date = DEFAULT_DATE + datetime.timedelta(days=i)
big_dag.create_dagrun(
execution_date=execution_date,
state=State.RUNNING,
run_type=DagRunType.SCHEDULED,
)
for task in big_dag.tasks:
ti = TI(task=copy.copy(task), execution_date=execution_date)
tis.append(ti)
import pendulum
tss = []
for ti in tis:
for i in range(5):
tss.append(TaskReschedule(
task=ti.task,
execution_date=ti.execution_date,
try_number=1,
start_date=pendulum.now(),
end_date=pendulum.now(),
reschedule_date=pendulum.now()))
def chunks(lst, n):
"""Yield successive n-sized chunks from lst."""
for i in range(0, len(lst), n):
yield lst[i:i + n]
with create_session() as session:
for chunk in chunks([{"task_id": tr.task_id, "dag_id": tr.dag_id, "execution_date": tr.execution_date,
"try_number": tr.try_number, "start_date": tr.start_date, "end_date": tr.end_date,
"duration": tr.duration, "reschedule_date": tr.reschedule_date} for tr in tss], 10000):
session.execute(TaskReschedule.__table__.insert(), chunk)
session.commit()
def test_clear_large_dag_with_task_reschedule(big_dag):
import time
with create_session() as session:
def count_task_reschedule():
return (
session.query(TaskReschedule)
.filter(
TaskReschedule.dag_id == big_dag.dag_id,
TaskReschedule.try_number == 1,
)
.count()
)
assert count_task_reschedule() == 1000 * 10 * 5
qry = session.query(TI).filter(TI.dag_id == big_dag.dag_id).all()
start = time.time()
clear_task_instances(qry, session, dag=big_dag)
end = time.time()
print(f"clear_task_instances took {end - start}")
assert count_task_reschedule() == 0
session.rollback()
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r570160429
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +170,23 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
Review comment:
Should we use `TaskInstanceKey` and:
https://github.com/apache/airflow/blob/2a3960f359b1fcd5eed64fff472c4e463da74beb/airflow/models/taskinstance.py#L1987
To avoid code duplication?
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +170,23 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
Review comment:
Should we apend `tis` and then use:
https://github.com/apache/airflow/blob/2a3960f359b1fcd5eed64fff472c4e463da74beb/airflow/models/taskinstance.py#L1987
to avoid code duplication?
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +170,23 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
Review comment:
Should we append `tis` and then use:
https://github.com/apache/airflow/blob/2a3960f359b1fcd5eed64fff472c4e463da74beb/airflow/models/taskinstance.py#L1987
to avoid code duplication?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] yuqian90 commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
yuqian90 commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r572006019
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +167,23 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
+
+ if tr_filter:
# Clear all reschedules related to the ti to clear
- session.query(TR).filter(
- TR.dag_id == ti.dag_id,
- TR.task_id == ti.task_id,
- TR.execution_date == ti.execution_date,
- TR.try_number == ti.try_number,
- ).delete()
+ delete_qry = TR.__table__.delete().where(
+ or_(
+ and_(
+ TR.dag_id == dag_id,
+ TR.task_id == task_id,
+ TR.execution_date == execution_date,
+ TR.try_number == try_number,
+ )
+ for dag_id, task_id, execution_date, try_number in tr_filter
Review comment:
Thanks. That's a great idea. I looked into the optimization in `filter_for_tis`. It works by extracting the common `dag_id` and `execution_date`, and then filter the `task_id` with an `IN` statement. This greatly reduces the size of the query, but it only does so if both `dag_id` and `execution_date` are equal for all tis.
For the case of slow `clear_task_instances`, the large number of tis is usually produced by the combination of a single dag_id over multiple `execution_date` for a large dag containing many task_ids. So it's slightly more complicated because there can be multiple `dag_id` and `execution_date` combinations (as well as `try_number`).
So I updated the logic to be more generic. It now uses a nested dictionary to produce the hierarchical filter. Some profiling showed great speed improvement (something like 40 to 50 times faster) compared to the first iteration. So the overral performance should now be 400 times faster than the original for loop deletion.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
kaxil commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r571463286
##########
File path: airflow/models/taskinstance.py
##########
@@ -145,7 +145,11 @@ def clear_task_instances(
:param activate_dag_runs: flag to check for active dag run
:param dag: DAG object
"""
+ from sqlalchemy import and_, or_
Review comment:
actually it is already at the top of file so not needed
```suggestion
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] yuqian90 commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
yuqian90 commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r572014808
##########
File path: airflow/models/taskinstance.py
##########
@@ -172,17 +173,34 @@ def clear_task_instances(
if tr_filter:
# Clear all reschedules related to the ti to clear
- delete_qry = TR.__table__.delete().where(
- or_(
- and_(
- TR.dag_id == dag_id,
- TR.task_id == task_id,
- TR.execution_date == execution_date,
- TR.try_number == try_number,
- )
- for dag_id, task_id, execution_date, try_number in tr_filter
+
+ # This is an optimization for the common case where all tis are for a small number
+ # of dag_id, execution_date and try_number. Use a nested dict of dag_id,
+ # execution_date, try_number and task_id to construct the where clause in a
+ # hierarchical manner. This speeds up the delete statement by more than 40x for
+ # large number of tis (50k+).
+ task_id_by_key = defaultdict(lambda: defaultdict(lambda: defaultdict(set)))
+ for dag_id, task_id, execution_date, try_number in tr_filter:
+ task_id_by_key[dag_id][execution_date][try_number].add(task_id)
Review comment:
Oh. Thanks @kaxil you are exactly right! Updated accordingly.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
kaxil commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r571459258
##########
File path: airflow/models/taskinstance.py
##########
@@ -145,7 +145,11 @@ def clear_task_instances(
:param activate_dag_runs: flag to check for active dag run
:param dag: DAG object
"""
+ from sqlalchemy import and_, or_
+
job_ids = []
+
Review comment:
```suggestion
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
kaxil commented on pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#issuecomment-774514107
Pushed https://github.com/apache/airflow/pull/14048/commits/99a77f3e368f5ea798deb5b97b30f14dbc298c2e -- Let's wait for the CI, will merge after that, thanks @yuqian90
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r570160429
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +170,23 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
Review comment:
Should we use `TaskInstanceKey` and:
https://github.com/apache/airflow/blob/2a3960f359b1fcd5eed64fff472c4e463da74beb/airflow/models/taskinstance.py#L1987
To avoid code duplication?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r570160429
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +170,23 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
Review comment:
Should we apend `tis` and then use:
https://github.com/apache/airflow/blob/2a3960f359b1fcd5eed64fff472c4e463da74beb/airflow/models/taskinstance.py#L1987
to avoid code duplication?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil merged pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
kaxil merged pull request #14048:
URL: https://github.com/apache/airflow/pull/14048
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] yuqian90 commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
yuqian90 commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r570123055
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +170,22 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
- # Clear all reschedules related to the ti to clear
- session.query(TR).filter(
- TR.dag_id == ti.dag_id,
- TR.task_id == ti.task_id,
- TR.execution_date == ti.execution_date,
- TR.try_number == ti.try_number,
- ).delete()
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
+
+ # Clear all reschedules related to the ti to clear
+ tr_qry = session.query(TR).filter(
+ or_(
+ and_(
+ TR.dag_id == dag_id,
+ TR.task_id == task_id,
+ TR.execution_date == execution_date,
+ TR.try_number == try_number,
+ )
+ for dag_id, task_id, execution_date, try_number in tr_filter
Review comment:
Thanks. That's a good point. We should defend against this. I updated according to your suggestion to not delete anything if `tr_filter` is empty.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
kaxil commented on pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#issuecomment-774513074
Minor suggestions added, LGTM otherwise
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] yuqian90 commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
yuqian90 commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r570123055
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +170,22 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
- # Clear all reschedules related to the ti to clear
- session.query(TR).filter(
- TR.dag_id == ti.dag_id,
- TR.task_id == ti.task_id,
- TR.execution_date == ti.execution_date,
- TR.try_number == ti.try_number,
- ).delete()
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
+
+ # Clear all reschedules related to the ti to clear
+ tr_qry = session.query(TR).filter(
+ or_(
+ and_(
+ TR.dag_id == dag_id,
+ TR.task_id == task_id,
+ TR.execution_date == execution_date,
+ TR.try_number == try_number,
+ )
+ for dag_id, task_id, execution_date, try_number in tr_filter
Review comment:
Thanks. That's a good point. We should defend against this. I updated according to your suggestion to not delete anything if `tr_filter` is empty.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] yuqian90 commented on pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
yuqian90 commented on pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#issuecomment-773213408
Thanks @potiuk and @mik-laj for the review. I updated the PR according to your comments.
Also, I noticed that doing `TR.__table__.delete().where(...)` seems to be slightly faster than what I was doing with `session.query(...).filter(...).delete()` by a small margin. So I changed the PR to use it instead.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
kaxil commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r572011423
##########
File path: airflow/models/taskinstance.py
##########
@@ -172,17 +173,34 @@ def clear_task_instances(
if tr_filter:
# Clear all reschedules related to the ti to clear
- delete_qry = TR.__table__.delete().where(
- or_(
- and_(
- TR.dag_id == dag_id,
- TR.task_id == task_id,
- TR.execution_date == execution_date,
- TR.try_number == try_number,
- )
- for dag_id, task_id, execution_date, try_number in tr_filter
+
+ # This is an optimization for the common case where all tis are for a small number
+ # of dag_id, execution_date and try_number. Use a nested dict of dag_id,
+ # execution_date, try_number and task_id to construct the where clause in a
+ # hierarchical manner. This speeds up the delete statement by more than 40x for
+ # large number of tis (50k+).
+ task_id_by_key = defaultdict(lambda: defaultdict(lambda: defaultdict(set)))
+ for dag_id, task_id, execution_date, try_number in tr_filter:
+ task_id_by_key[dag_id][execution_date][try_number].add(task_id)
Review comment:
Can you post the code that calculates times takes to see 40x speed
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r569832257
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +170,22 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
- # Clear all reschedules related to the ti to clear
- session.query(TR).filter(
- TR.dag_id == ti.dag_id,
- TR.task_id == ti.task_id,
- TR.execution_date == ti.execution_date,
- TR.try_number == ti.try_number,
- ).delete()
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
+
+ # Clear all reschedules related to the ti to clear
+ tr_qry = session.query(TR).filter(
+ or_(
+ and_(
+ TR.dag_id == dag_id,
+ TR.task_id == task_id,
+ TR.execution_date == execution_date,
+ TR.try_number == try_number,
+ )
+ for dag_id, task_id, execution_date, try_number in tr_filter
Review comment:
Is `tr_filter` always contain at least one element? I'm afraid that if `tr_filter` is empty, too many objects in the database will be deleted because ORM won't generate SQL query without any filters.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
kaxil commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r572010717
##########
File path: airflow/models/taskinstance.py
##########
@@ -172,17 +173,34 @@ def clear_task_instances(
if tr_filter:
# Clear all reschedules related to the ti to clear
- delete_qry = TR.__table__.delete().where(
- or_(
- and_(
- TR.dag_id == dag_id,
- TR.task_id == task_id,
- TR.execution_date == execution_date,
- TR.try_number == try_number,
- )
- for dag_id, task_id, execution_date, try_number in tr_filter
+
+ # This is an optimization for the common case where all tis are for a small number
+ # of dag_id, execution_date and try_number. Use a nested dict of dag_id,
+ # execution_date, try_number and task_id to construct the where clause in a
+ # hierarchical manner. This speeds up the delete statement by more than 40x for
+ # large number of tis (50k+).
+ task_id_by_key = defaultdict(lambda: defaultdict(lambda: defaultdict(set)))
+ for dag_id, task_id, execution_date, try_number in tr_filter:
+ task_id_by_key[dag_id][execution_date][try_number].add(task_id)
Review comment:
We can move this in the loop above instead of looping across it again, right?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] yuqian90 commented on a change in pull request #14048: Speed up clear_task_instances by doing a single sql delete for TaskReschedule
Posted by GitBox <gi...@apache.org>.
yuqian90 commented on a change in pull request #14048:
URL: https://github.com/apache/airflow/pull/14048#discussion_r571452659
##########
File path: airflow/models/taskinstance.py
##########
@@ -166,13 +170,23 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
ti.state = State.NONE
session.merge(ti)
+
+ tr_filter.append((ti.dag_id, ti.task_id, ti.execution_date, ti.try_number))
Review comment:
Hi @turbaszek. I looked into `filter_for_tis()`. It specifically creates a filter for `TaskInstance`. In this case, we need to filter for `TaskReschedule` instead. Also, it does not include a condition for `try_number`, which is necessary to get `TaskReschedule` right.
```python
return or_(
and_(
TaskInstance.dag_id == ti.dag_id,
TaskInstance.task_id == ti.task_id,
TaskInstance.execution_date == ti.execution_date,
)
for ti in tis
)
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org