Speed up clear_task_instances by doing a single sql delete for TaskReschedule

[yuqian90]

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.

        # 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 OR 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.

[potiuk]

Looks good to me. @kaxil - can you take a look on that one as well? It looks like nice thing to cherry-pick to 2.0.1, very localized and seems to improve the case that is rather useful.

[github-actions]

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.

[yuqian90]

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.

[kaxil]

Minor suggestions added, LGTM otherwise

[kaxil]

Pushed 99a77f3 -- Let's wait for the CI, will merge after that, thanks @yuqian90

[yuqian90]

Followed @ashb 's suggestion and speed it up even further. Please take another look.

[yuqian90]

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.

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()