Skip to content

When clearing task instances try to get associated DAGs from database #29065

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
potiuk merged 4 commits into from
Apr 14, 2023
Merged

When clearing task instances try to get associated DAGs from database #29065

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e8c6a89
When clearing task instances try to get associated DAGs from database.
sean-rose Jan 20, 2023
4efc224
Use session fixture for new `test_clear_task_instances_without_dag_pa…
sean-rose Jan 23, 2023
3958df0
Use session fixture for new `test_clear_task_instances_in_multiple_da…
sean-rose Jan 23, 2023
c5c0591
Merge branch 'main' into fix-recursive-clear
eladkal Apr 14, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions airflow/models/taskinstance.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,7 @@
)
from airflow.listeners.listener import get_listener_manager
from airflow.models.base import Base, StringID
from airflow.models.dagbag import DagBag
from airflow.models.log import Log
from airflow.models.mappedoperator import MappedOperator
from airflow.models.param import process_params
Expand Down Expand Up @@ -203,6 +204,7 @@ def clear_task_instances(
task_id_by_key: dict[str, dict[str, dict[int, dict[int, set[str]]]]] = defaultdict(
lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(set)))
)
dag_bag = DagBag(read_dags_from_db=True)

@uranusjr uranusjr Feb 8, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder why dag is not a required argument to begin with. It really should be.

@dimberman dimberman Feb 13, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ti_dag = dag if dag and dag.dag_id == ti.dag_id else dag_bag.get_dag(ti.dag_id, session=session)

Looks like dag is optional because if there is no DAG we can just get the DAG from the TaskInstance

@uranusjr uranusjr Feb 13, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With this change yes, but without this PR dag=None would have a very weird behaviour and not particularly useful. This makes me think the signature may be considered a bug in the first place and fixed without introducing DagBag here.

@dstandish dstandish May 1, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Blast from a past, but I just had reason to encounter this headscratcher of a function also. I think it can be that we can deal with many different dags in this function, because e.g. there is no restriction that the TIs passed in must be from the same dag. And I think this gets used in real life when you have external task sensors and you clear the upstream task. Too stinkin complicated.

for ti in tis:
if ti.state == TaskInstanceState.RUNNING:
if ti.job_id:
Expand All @@ -211,15 +213,16 @@ def clear_task_instances(
ti.state = TaskInstanceState.RESTARTING
job_ids.append(ti.job_id)
else:
ti_dag = dag if dag and dag.dag_id == ti.dag_id else dag_bag.get_dag(ti.dag_id, session=session)
task_id = ti.task_id
if dag and dag.has_task(task_id):
task = dag.get_task(task_id)
if ti_dag and ti_dag.has_task(task_id):
task = ti_dag.get_task(task_id)
ti.refresh_from_task(task)
task_retries = task.retries
ti.max_tries = ti.try_number + task_retries - 1
else:
# Ignore errors when updating max_tries if dag is None or
# task not found in dag since database records could be
# Ignore errors when updating max_tries if the DAG or
# task are not found since database records could be
# outdated. We make max_tries the maximum value of its
# original max_tries or the last attempted try number.
ti.max_tries = max(ti.max_tries, ti.prev_attempted_tries)
Expand Down
97 changes: 94 additions & 3 deletions tests/models/test_cleartasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

from airflow import settings
from airflow.models import DAG, TaskInstance as TI, TaskReschedule, clear_task_instances
from airflow.models.serialized_dag import SerializedDagModel
from airflow.operators.empty import EmptyOperator
from airflow.sensors.python import PythonSensor
from airflow.utils.session import create_session
Expand Down Expand Up @@ -202,9 +203,9 @@ def test_clear_task_instances_without_task(self, dag_maker):
# but it works for our case because we specifically constructed test DAGS
# in the way that those two sort methods are equivalent
qry = session.query(TI).filter(TI.dag_id == dag.dag_id).order_by(TI.task_id).all()
clear_task_instances(qry, session)
clear_task_instances(qry, session, dag=dag)
Comment thread
ephraimbuddy marked this conversation as resolved.

# When dag is None, max_tries will be maximum of original max_tries or try_number.
# When no task is found, max_tries will be maximum of original max_tries or try_number.
ti0.refresh_from_db()
ti1.refresh_from_db()
# Next try to run will be try 2
Expand All @@ -214,6 +215,7 @@ def test_clear_task_instances_without_task(self, dag_maker):
assert ti1.max_tries == 2

def test_clear_task_instances_without_dag(self, dag_maker):
# Don't write DAG to the database, so no DAG is found by clear_task_instances().
with dag_maker(
"test_clear_task_instances_without_dag",
start_date=DEFAULT_DATE,
Expand Down Expand Up @@ -242,7 +244,7 @@ def test_clear_task_instances_without_dag(self, dag_maker):
qry = session.query(TI).filter(TI.dag_id == dag.dag_id).order_by(TI.task_id).all()
clear_task_instances(qry, session)

# When dag is None, max_tries will be maximum of original max_tries or try_number.
# When no DAG is found, max_tries will be maximum of original max_tries or try_number.
ti0.refresh_from_db()
ti1.refresh_from_db()
# Next try to run will be try 2
Expand All @@ -251,6 +253,95 @@ def test_clear_task_instances_without_dag(self, dag_maker):
assert ti1.try_number == 2
assert ti1.max_tries == 2

def test_clear_task_instances_without_dag_param(self, dag_maker, session):
with dag_maker(
"test_clear_task_instances_without_dag_param",
start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE + datetime.timedelta(days=10),
session=session,
) as dag:
task0 = EmptyOperator(task_id="task0")
task1 = EmptyOperator(task_id="task1", retries=2)

# Write DAG to the database so it can be found by clear_task_instances().
SerializedDagModel.write_dag(dag, session=session)

dr = dag_maker.create_dagrun(
state=State.RUNNING,
run_type=DagRunType.SCHEDULED,
)

ti0, ti1 = sorted(dr.task_instances, key=lambda ti: ti.task_id)
ti0.refresh_from_task(task0)
ti1.refresh_from_task(task1)

ti0.run(session=session)
ti1.run(session=session)

# we use order_by(task_id) here because for the test DAG structure of ours
# this is equivalent to topological sort. It would not work in general case
# but it works for our case because we specifically constructed test DAGS
# in the way that those two sort methods are equivalent
qry = session.query(TI).filter(TI.dag_id == dag.dag_id).order_by(TI.task_id).all()
clear_task_instances(qry, session)

ti0.refresh_from_db(session=session)
ti1.refresh_from_db(session=session)
# Next try to run will be try 2
assert ti0.try_number == 2
assert ti0.max_tries == 1
assert ti1.try_number == 2
assert ti1.max_tries == 3

def test_clear_task_instances_in_multiple_dags(self, dag_maker, session):
with dag_maker(
"test_clear_task_instances_in_multiple_dags0",
start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE + datetime.timedelta(days=10),
session=session,
) as dag0:
task0 = EmptyOperator(task_id="task0")

dr0 = dag_maker.create_dagrun(
state=State.RUNNING,
run_type=DagRunType.SCHEDULED,
)

with dag_maker(
"test_clear_task_instances_in_multiple_dags1",
start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE + datetime.timedelta(days=10),
session=session,
) as dag1:
task1 = EmptyOperator(task_id="task1", retries=2)

# Write secondary DAG to the database so it can be found by clear_task_instances().
SerializedDagModel.write_dag(dag1, session=session)

dr1 = dag_maker.create_dagrun(
state=State.RUNNING,
run_type=DagRunType.SCHEDULED,
)

ti0 = dr0.task_instances[0]
ti1 = dr1.task_instances[0]
ti0.refresh_from_task(task0)
ti1.refresh_from_task(task1)

ti0.run(session=session)
ti1.run(session=session)

qry = session.query(TI).filter(TI.dag_id.in_((dag0.dag_id, dag1.dag_id))).all()
clear_task_instances(qry, session, dag=dag0)

ti0.refresh_from_db(session=session)
ti1.refresh_from_db(session=session)
# Next try to run will be try 2
assert ti0.try_number == 2
assert ti0.max_tries == 1
assert ti1.try_number == 2
assert ti1.max_tries == 3

def test_clear_task_instances_with_task_reschedule(self, dag_maker):
"""Test that TaskReschedules are deleted correctly when TaskInstances are cleared"""

Expand Down