Repository: apache/airflow
Pull request #49752: Mark DagRun as success when no teardown tasks are running
Author: login: renzepost
Labels: name: area:API
description: Airflow's REST/HTTP API
Created: 2025-04-24T20:55:58Z | Merged: 2025-04-30T11:05:11Z

<!--
 Licensed to the Apache Software Foundation (ASF) under one
 or more contributor license agreements.  See the NOTICE file
 distributed with this work for additional information
 regarding copyright ownership.  The ASF licenses this file
 to you under the Apache License, Version 2.0 (the
 "License"); you may not use this file except in compliance
 with the License.  You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

 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
 KIND, either express or implied.  See the License for the
 specific language governing permissions and limitations
 under the License.
 -->

<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.

Feel free to ping committers for the review!

In case of an existing issue, reference it using one of the following:

closes: #ISSUE
related: #ISSUE

How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
closes: #46937 


<!-- Please keep an empty line above the dashes. -->
---
When marking a dag run as success, first check if any teardown tasks are in state `running`, `deferred` or `up_for_reschedule` and if so, don't mark the dag run as success (because the teardowns need to complete). If all teardown tasks are in any other state, a dag run should be marked as success.
 
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).



Diff:
diff --git a/airflow-core/src/airflow/api/common/mark_tasks.py b/airflow-core/src/airflow/api/common/mark_tasks.py
index c957a5cd53ab3..ad8c7fe4928cd 100644
--- a/airflow-core/src/airflow/api/common/mark_tasks.py
+++ b/airflow-core/src/airflow/api/common/mark_tasks.py
@@ -215,17 +215,31 @@ def set_dag_run_state_to_success(
     if not run_id:
         raise ValueError(f"Invalid dag_run_id: {run_id}")
 
-    # Mark all task instances of the dag run to success - except for teardown as they need to complete work.
+    # Mark all task instances of the dag run to success - except for unfinished teardown as they need to complete work.
     normal_tasks = [task for task in dag.tasks if not task.is_teardown]
+    teardown_tasks = [task for task in dag.tasks if task.is_teardown]
+    unfinished_teardown_task_ids = set(
+        session.scalars(
+            select(TaskInstance.task_id).where(
+                TaskInstance.dag_id == dag.dag_id,
+                TaskInstance.run_id == run_id,
+                TaskInstance.task_id.in_([task.task_id for task in teardown_tasks]),
+                or_(TaskInstance.state.is_(None), TaskInstance.state.in_(State.unfinished)),
+            )
+        ).all()
+    )
 
-    # Mark the dag run to success.
-    if commit and len(normal_tasks) == len(dag.tasks):
+    # Mark the dag run to success if there are no unfinished teardown tasks.
+    if commit and len(unfinished_teardown_task_ids) == 0:
         _set_dag_run_state(dag.dag_id, run_id, DagRunState.SUCCESS, session)
 
-    for task in normal_tasks:
+    tasks_to_mark_success = normal_tasks + [
+        task for task in teardown_tasks if task.task_id not in unfinished_teardown_task_ids
+    ]
+    for task in tasks_to_mark_success:
         task.dag = dag
     return set_state(
-        tasks=normal_tasks,
+        tasks=tasks_to_mark_success,
         run_id=run_id,
         state=TaskInstanceState.SUCCESS,
         commit=commit,
diff --git a/airflow-core/tests/unit/api/common/test_mark_tasks.py b/airflow-core/tests/unit/api/common/test_mark_tasks.py
index 59fca1f4caaf8..35dfe7b58793b 100644
--- a/airflow-core/tests/unit/api/common/test_mark_tasks.py
+++ b/airflow-core/tests/unit/api/common/test_mark_tasks.py
@@ -19,10 +19,12 @@
 from typing import TYPE_CHECKING
 
 import pytest
+from sqlalchemy import select
 
 from airflow.api.common.mark_tasks import set_dag_run_state_to_failed, set_dag_run_state_to_success
+from airflow.models.dagrun import DagRun
 from airflow.providers.standard.operators.empty import EmptyOperator
-from airflow.utils.state import TaskInstanceState
+from airflow.utils.state import DagRunState, State, TaskInstanceState
 
 if TYPE_CHECKING:
     from airflow.models.taskinstance import TaskInstance
@@ -54,23 +56,63 @@ def test_set_dag_run_state_to_failed(dag_maker: DagMaker):
     assert "teardown" not in task_dict
 
 
-def test_set_dag_run_state_to_success(dag_maker: DagMaker):
+@pytest.mark.parametrize(
+    "unfinished_state", sorted([state for state in State.unfinished if state is not None])
+)
+def test_set_dag_run_state_to_success_unfinished_teardown(dag_maker: DagMaker, unfinished_state):
     with dag_maker("TEST_DAG_1"):
         with EmptyOperator(task_id="teardown").as_teardown():
             EmptyOperator(task_id="running")
             EmptyOperator(task_id="pending")
+
     dr = dag_maker.create_dagrun()
     for ti in dr.get_task_instances():
         if ti.task_id == "running":
             ti.set_state(TaskInstanceState.RUNNING)
+        if ti.task_id == "teardown":
+            ti.set_state(unfinished_state)
+
     dag_maker.session.flush()
     assert dr.dag
+    assert dr.state == DagRunState.RUNNING
 
     updated_tis: list[TaskInstance] = set_dag_run_state_to_success(
         dag=dr.dag, run_id=dr.run_id, commit=True, session=dag_maker.session
     )
+    run = dag_maker.session.scalar(select(DagRun).filter_by(dag_id=dr.dag_id, run_id=dr.run_id))
+    assert run.state != DagRunState.SUCCESS
     assert len(updated_tis) == 2
     task_dict = {ti.task_id: ti for ti in updated_tis}
     assert task_dict["running"].state == TaskInstanceState.SUCCESS
     assert task_dict["pending"].state == TaskInstanceState.SUCCESS
     assert "teardown" not in task_dict
+
+
+@pytest.mark.parametrize("finished_state", sorted(list(State.finished)))
+def test_set_dag_run_state_to_success_finished_teardown(dag_maker: DagMaker, finished_state):
+    with dag_maker("TEST_DAG_1"):
+        with EmptyOperator(task_id="teardown").as_teardown():
+            EmptyOperator(task_id="failed")
+    dr = dag_maker.create_dagrun()
+    for ti in dr.get_task_instances():
+        if ti.task_id == "failed":
+            ti.set_state(TaskInstanceState.FAILED)
+        if ti.task_id == "teardown":
+            ti.set_state(finished_state)
+    dag_maker.session.flush()
+    dr.set_state(DagRunState.FAILED)
+    assert dr.dag
+
+    updated_tis: list[TaskInstance] = set_dag_run_state_to_success(
+        dag=dr.dag, run_id=dr.run_id, commit=True, session=dag_maker.session
+    )
+    run = dag_maker.session.scalar(select(DagRun).filter_by(dag_id=dr.dag_id, run_id=dr.run_id))
+    assert run.state == DagRunState.SUCCESS
+    if finished_state == TaskInstanceState.SUCCESS:
+        assert len(updated_tis) == 1
+    else:
+        assert len(updated_tis) == 2
+    task_dict = {ti.task_id: ti for ti in updated_tis}
+    assert task_dict["failed"].state == TaskInstanceState.SUCCESS
+    if finished_state != TaskInstanceState.SUCCESS:
+        assert task_dict["teardown"].state == TaskInstanceState.SUCCESS


Reviews:
user: login: pierrejeambrun
body: Do we also need (maybe in another PR) to update the mark as failed ? To also not early fail teardown tasks?
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-28T09:30:51Z
user: login: renzepost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T15:36:04Z
user: login: pierrejeambrun
body: Looks good ty.
state: DISMISSED
author_association: MEMBER
submitted_at: 2025-04-29T08:20:53Z
user: login: pierrejeambrun
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-29T08:32:19Z
user: login: uranusjr
body: Some minor suggestions; this looks good to me too.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-29T09:17:44Z
user: login: renzepost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T10:16:25Z
user: login: renzepost
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T10:16:56Z
user: login: pierrejeambrun
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-30T10:03:59Z

Review comments:
diff_hunk: @@ -215,17 +215,35 @@ def set_dag_run_state_to_success(
     if not run_id:
         raise ValueError(f"Invalid dag_run_id: {run_id}")
 
-    # Mark all task instances of the dag run to success - except for teardown as they need to complete work.
+    # Mark all task instances of the dag run to success - except for running teardown as they need to complete work.
     normal_tasks = [task for task in dag.tasks if not task.is_teardown]
+    teardown_tasks = [task for task in dag.tasks if task.is_teardown]
+    running_states = (
+        TaskInstanceState.RUNNING,
+        TaskInstanceState.DEFERRED,
+        TaskInstanceState.UP_FOR_RESCHEDULE,
+    )
path: airflow-core/src/airflow/api/common/mark_tasks.py
user: login: pierrejeambrun
body: Should we instead use `State.unfinished` ? I know we do that in the `mark_as_failed` function, but it seems that `TaskInstanceState.RESTARTING`, or `TaskInstanceState.UP_FOR_RETRY` etc... shouldn't be marked as success too early.
created_at: 2025-04-28T09:30:49Z
author_association: MEMBER
diff_hunk: @@ -215,17 +215,35 @@ def set_dag_run_state_to_success(
     if not run_id:
         raise ValueError(f"Invalid dag_run_id: {run_id}")
 
-    # Mark all task instances of the dag run to success - except for teardown as they need to complete work.
+    # Mark all task instances of the dag run to success - except for running teardown as they need to complete work.
     normal_tasks = [task for task in dag.tasks if not task.is_teardown]
+    teardown_tasks = [task for task in dag.tasks if task.is_teardown]
+    running_states = (
+        TaskInstanceState.RUNNING,
+        TaskInstanceState.DEFERRED,
+        TaskInstanceState.UP_FOR_RESCHEDULE,
+    )
path: airflow-core/src/airflow/api/common/mark_tasks.py
user: login: renzepost
body: Wasn't aware of `State.unfinished`, agreed that's much better! Updated. Also changed the unit tests to cover the possible states.
created_at: 2025-04-28T15:36:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -54,23 +56,61 @@ def test_set_dag_run_state_to_failed(dag_maker: DagMaker):
     assert "teardown" not in task_dict
 
 
-def test_set_dag_run_state_to_success(dag_maker: DagMaker):
+@pytest.mark.parametrize(
+    "unfinished_state", sorted([state for state in State.unfinished if state is not None])
+)
+def test_set_dag_run_state_to_success_unfinished_teardown(dag_maker: DagMaker, unfinished_state):
     with dag_maker("TEST_DAG_1"):
         with EmptyOperator(task_id="teardown").as_teardown():
             EmptyOperator(task_id="running")
             EmptyOperator(task_id="pending")
+
     dr = dag_maker.create_dagrun()
     for ti in dr.get_task_instances():
         if ti.task_id == "running":
             ti.set_state(TaskInstanceState.RUNNING)
+        if ti.task_id == "teardown":
+            ti.set_state(unfinished_state)
+
     dag_maker.session.flush()
     assert dr.dag
 
     updated_tis: list[TaskInstance] = set_dag_run_state_to_success(
         dag=dr.dag, run_id=dr.run_id, commit=True, session=dag_maker.session
     )
+    run = dag_maker.session.scalar(select(DagRun).filter_by(dag_id=dr.dag_id, run_id=dr.run_id))
+    assert run.state != DagRunState.SUCCESS
     assert len(updated_tis) == 2
     task_dict = {ti.task_id: ti for ti in updated_tis}
     assert task_dict["running"].state == TaskInstanceState.SUCCESS
     assert task_dict["pending"].state == TaskInstanceState.SUCCESS
     assert "teardown" not in task_dict
+
+
+@pytest.mark.parametrize(
+    "finished_state", sorted([state for state in State.finished if state != TaskInstanceState.SUCCESS])
+)
path: airflow-core/tests/unit/api/common/test_mark_tasks.py
user: login: pierrejeambrun
body: I think we need to keep success there too. (that's the use case mentioned in the issue)

Also I think we should assert dr state is `Failed` before marking it as success
created_at: 2025-04-29T08:32:15Z
author_association: MEMBER
diff_hunk: @@ -215,17 +215,30 @@ def set_dag_run_state_to_success(
     if not run_id:
         raise ValueError(f"Invalid dag_run_id: {run_id}")
 
-    # Mark all task instances of the dag run to success - except for teardown as they need to complete work.
+    # Mark all task instances of the dag run to success - except for unfinished teardown as they need to complete work.
     normal_tasks = [task for task in dag.tasks if not task.is_teardown]
+    teardown_tasks = [task for task in dag.tasks if task.is_teardown]
+    unfinished_teardown_tis: list[TaskInstance] = session.scalars(
+        select(TaskInstance).where(
+            TaskInstance.dag_id == dag.dag_id,
+            TaskInstance.run_id == run_id,
+            TaskInstance.task_id.in_([task.task_id for task in teardown_tasks]),
+            or_(TaskInstance.state.is_(None), TaskInstance.state.in_(State.unfinished)),
+        )
+    ).all()
path: airflow-core/src/airflow/api/common/mark_tasks.py
user: login: uranusjr
body: We only really needs the task_id here, so let’s just select that field instead of the entire TaskInstance object. Also putting the task_ids in a set (instead of list) would be a good idea.
created_at: 2025-04-29T09:15:27Z
author_association: MEMBER
diff_hunk: @@ -54,23 +56,61 @@ def test_set_dag_run_state_to_failed(dag_maker: DagMaker):
     assert "teardown" not in task_dict
 
 
-def test_set_dag_run_state_to_success(dag_maker: DagMaker):
+@pytest.mark.parametrize(
+    "unfinished_state", sorted([state for state in State.unfinished if state is not None])
+)
+def test_set_dag_run_state_to_success_unfinished_teardown(dag_maker: DagMaker, unfinished_state):
     with dag_maker("TEST_DAG_1"):
         with EmptyOperator(task_id="teardown").as_teardown():
             EmptyOperator(task_id="running")
             EmptyOperator(task_id="pending")
+
     dr = dag_maker.create_dagrun()
     for ti in dr.get_task_instances():
         if ti.task_id == "running":
             ti.set_state(TaskInstanceState.RUNNING)
+        if ti.task_id == "teardown":
+            ti.set_state(unfinished_state)
+
     dag_maker.session.flush()
     assert dr.dag
 
     updated_tis: list[TaskInstance] = set_dag_run_state_to_success(
         dag=dr.dag, run_id=dr.run_id, commit=True, session=dag_maker.session
     )
+    run = dag_maker.session.scalar(select(DagRun).filter_by(dag_id=dr.dag_id, run_id=dr.run_id))
+    assert run.state != DagRunState.SUCCESS
     assert len(updated_tis) == 2
     task_dict = {ti.task_id: ti for ti in updated_tis}
     assert task_dict["running"].state == TaskInstanceState.SUCCESS
     assert task_dict["pending"].state == TaskInstanceState.SUCCESS
     assert "teardown" not in task_dict
+
+
+@pytest.mark.parametrize(
+    "finished_state", sorted([state for state in State.finished if state != TaskInstanceState.SUCCESS])
+)
path: airflow-core/tests/unit/api/common/test_mark_tasks.py
user: login: renzepost
body: Updated, now success is also tested. Also changed the scenario of the last unit test a bit, the dagrun is set to `Failed` before marking it as success. I think it makes more sense now.
created_at: 2025-04-29T10:16:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -215,17 +215,30 @@ def set_dag_run_state_to_success(
     if not run_id:
         raise ValueError(f"Invalid dag_run_id: {run_id}")
 
-    # Mark all task instances of the dag run to success - except for teardown as they need to complete work.
+    # Mark all task instances of the dag run to success - except for unfinished teardown as they need to complete work.
     normal_tasks = [task for task in dag.tasks if not task.is_teardown]
+    teardown_tasks = [task for task in dag.tasks if task.is_teardown]
+    unfinished_teardown_tis: list[TaskInstance] = session.scalars(
+        select(TaskInstance).where(
+            TaskInstance.dag_id == dag.dag_id,
+            TaskInstance.run_id == run_id,
+            TaskInstance.task_id.in_([task.task_id for task in teardown_tasks]),
+            or_(TaskInstance.state.is_(None), TaskInstance.state.in_(State.unfinished)),
+        )
+    ).all()
path: airflow-core/src/airflow/api/common/mark_tasks.py
user: login: renzepost
body: Done! Thanks for the feedback.
created_at: 2025-04-29T10:16:56Z
author_association: CONTRIBUTOR

PR comments:
user: login: renzepost
created_at: 2025-04-28T15:39:29Z
body: > Do we also need (maybe in another PR) to update the mark as failed ? To also not early fail teardown tasks?

I didn't investigate this thoroughly, but I think mark failed already works as expected. Or do you think the `running_states` there might not cover all the cases? I could create another PR if that's the case, or if I don't have time I could at least create another issue for someone else to pick up.
author_association: CONTRIBUTOR
user: login: pierrejeambrun
created_at: 2025-04-29T08:23:16Z
body: > Or do you think the running_states there might not cover all the cases?

I mean mark as failed is coming from another function, `set_dag_run_state_to_failed` that is apparently not handling `teardown` tasks (there's no special case for them), so basically they can be marked as `failed` early on, without letting them complete beforehand. But maybe that's expected 🤔 not sure
author_association: MEMBER
user: login: pierrejeambrun
created_at: 2025-04-29T08:30:05Z
body: Also this does not really implement the related issue report. ~~Maybe we need clarification on the expected behavior~~

edit: This change makes sense on its own but will not close the related issue. (or can we add a test for it, teardown need to be success, 1 task in the dag need to be failed, and we are marking as success the dagrun, which will remain wrongly in failed state).

I would be for merging this on its own, and create a separate PR for the related issue. (unless this one resolves it too)
author_association: MEMBER
user: login: pierrejeambrun
created_at: 2025-04-29T09:01:34Z
body: Can you rebase the branch, this will fix the CI issue.
author_association: MEMBER
user: login: renzepost
created_at: 2025-04-29T10:23:45Z
body: > edit: This change makes sense on its own but will not close the related issue. (or can we add a test for it, teardown need to be success, 1 task in the dag need to be failed, and we are marking as success the dagrun, which will remain wrongly in failed state).

The last unit test now covers this scenario with `TaskInstanceState.SUCCESS` parameter. There's a failed "normal" task, the dagrun is in failed state, the teardown will be in success state, and after marking the dagrun to success, the dagrun is indeed marked as success (and so is the failed "normal" task).
author_association: CONTRIBUTOR
user: login: pierrejeambrun
created_at: 2025-04-30T10:03:05Z
body: > The last unit test now covers this scenario with TaskInstanceState.SUCCESS parameter. There's a failed "normal" task, the dagrun is in failed state, the teardown will be in success state, and after marking the dagrun to success, the dagrun is indeed marked as success (and so is the failed "normal" task).


Great, then it's closing the related issue! 
author_association: MEMBER
user: login: github-actions[bot]
created_at: 2025-04-30T11:06:02Z
body: ### Backport successfully created: v3-0-test

<table>
                <tr>
                    <th>Status</th>
                    <th>Branch</th>
                    <th>Result</th>
                </tr>
                <tr>
                    <td>✅</td>
                    <td>v3-0-test</td>
                    <td><a href="https://github.com/apache/airflow/pull/50019"><img src="https://img.shields.io/badge/PR-50019-blue" alt="PR Link"></a></td>
                </tr>
            </table>
author_association: CONTRIBUTOR