This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch cherry-pick-69821-v3-3-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 5a43971acc19511ed6f21d0469644f92b1da0512 Author: Hemkumar Chheda <[email protected]> AuthorDate: Wed Aug 5 14:05:31 2026 +0530 Respect retries when a deferrable trigger ends a task with TaskFailedEvent (#69821) * Respect retries when a deferrable trigger ends a task with TaskFailedEvent When a deferrable operator's trigger yields a terminal TaskFailedEvent, the task was always marked failed and on_failure_callback ran, even with retries remaining. A worker-side failure with retries left instead goes up_for_retry and runs on_retry_callback. Route trigger-emitted failures through TaskInstance.handle_failure (the path the scheduler has used since #56586) so retry-eligibility is respected and the callback runs on_retry_callback vs on_failure_callback accordingly. closes: #69819 Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]> * Preserve retry history for trigger-ended deferred failures * Fix stale trigger callback comment --------- Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]> (cherry picked from commit dd1f84ab2034a14164e4d09310fba8fe1fa0395e) --- airflow-core/newsfragments/69821.bugfix.rst | 1 + airflow-core/src/airflow/models/trigger.py | 46 ++++++++++++++-- airflow-core/tests/unit/models/test_trigger.py | 76 +++++++++++++++++++++++++- 3 files changed, 115 insertions(+), 8 deletions(-) diff --git a/airflow-core/newsfragments/69821.bugfix.rst b/airflow-core/newsfragments/69821.bugfix.rst new file mode 100644 index 00000000000..9dca1eb7d0c --- /dev/null +++ b/airflow-core/newsfragments/69821.bugfix.rst @@ -0,0 +1 @@ +Deferrable tasks that fail via a trigger-emitted ``TaskFailedEvent`` now respect retries: if the task has retries remaining it goes ``up_for_retry`` and runs ``on_retry_callback``, instead of always failing terminally and running ``on_failure_callback``. diff --git a/airflow-core/src/airflow/models/trigger.py b/airflow-core/src/airflow/models/trigger.py index 597513a3d07..9d8c195e9c9 100644 --- a/airflow-core/src/airflow/models/trigger.py +++ b/airflow-core/src/airflow/models/trigger.py @@ -560,19 +560,41 @@ def _(event: BaseTaskEndEvent, *, task_instance: TaskInstance, session: Session) from airflow.callbacks.database_callback_sink import DatabaseCallbackSink from airflow.utils.state import TaskInstanceState - # Mark the task with terminal state and prevent it from resuming on worker + # Prevent the task from resuming on a worker. task_instance.trigger_id = None - task_instance.set_state(event.task_instance_state, session=session) + + callback_type = event.task_instance_state + should_retry = False + + if event.task_instance_state == TaskInstanceState.FAILED: + # Load the serialized task so retry eligibility matches the normal task path. + try: + from airflow.models.dagbag import DBDagBag + + dag = DBDagBag().get_dag_for_run(dag_run=task_instance.dag_run, session=session) + if dag is not None: + task_instance.task = dag.get_task(task_instance.task_id) + should_retry = task_instance.is_eligible_to_retry() + except Exception: + log.exception( + "Could not load task for %s; failing terminally without retry routing", task_instance + ) + if should_retry: + callback_type = TaskInstanceState.UP_FOR_RETRY def _submit_callback_if_necessary() -> None: - """Submit a callback request if the task state is SUCCESS or FAILED.""" - if event.task_instance_state in (TaskInstanceState.SUCCESS, TaskInstanceState.FAILED): + """Submit a callback request if the task state is SUCCESS, FAILED, or UP_FOR_RETRY.""" + if callback_type in ( + TaskInstanceState.SUCCESS, + TaskInstanceState.FAILED, + TaskInstanceState.UP_FOR_RETRY, + ): if task_instance.dag_model.relative_fileloc is None: raise RuntimeError("relative_fileloc should not be None for a finished task") request = TaskCallbackRequest( filepath=task_instance.dag_model.relative_fileloc, ti=task_instance, - task_callback_type=event.task_instance_state, + task_callback_type=callback_type, bundle_name=task_instance.dag_model.bundle_name, bundle_version=task_instance.dag_run.bundle_version, ) @@ -584,10 +606,22 @@ def _(event: BaseTaskEndEvent, *, task_instance: TaskInstance, session: Session) def _push_xcoms_if_necessary() -> None: """Pushes XComs to the database if they are provided.""" - if event.xcoms: + if event.xcoms and callback_type != TaskInstanceState.UP_FOR_RETRY: for key, value in event.xcoms.items(): task_instance.xcom_push(key=key, value=value) + # Send the callback before mutating task state so it reflects the retry-vs-terminal + # decision derived above. _submit_callback_if_necessary() + + if should_retry: + task_instance.end_date = timezone.utcnow() + task_instance.set_duration() + task_instance.clear_next_method_args() + task_instance.prepare_db_for_next_try(session) + task_instance.state = TaskInstanceState.UP_FOR_RETRY + else: + task_instance.set_state(event.task_instance_state, session=session) + _push_xcoms_if_necessary() session.flush() diff --git a/airflow-core/tests/unit/models/test_trigger.py b/airflow-core/tests/unit/models/test_trigger.py index 669d98a7704..2785c48c1dd 100644 --- a/airflow-core/tests/unit/models/test_trigger.py +++ b/airflow-core/tests/unit/models/test_trigger.py @@ -34,6 +34,7 @@ from airflow.jobs.triggerer_job_runner import TriggererJobRunner from airflow.models import TaskInstance, Trigger from airflow.models.asset import AssetEvent, AssetModel, AssetWatcherModel from airflow.models.callback import Callback, TriggererCallback +from airflow.models.taskinstancehistory import TaskInstanceHistory from airflow.models.xcom import XComModel from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk.definitions.callback import AsyncCallback @@ -48,7 +49,7 @@ from airflow.triggers.base import ( TriggerEvent, ) from airflow.utils.session import create_session -from airflow.utils.state import State +from airflow.utils.state import State, TaskInstanceState from tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.config import conf_vars @@ -296,11 +297,15 @@ def test_submit_event_task_end(mock_utcnow, session, create_task_instance, event # Make a trigger trigger = Trigger(classpath="does.not.matter", kwargs={}) session.add(trigger) - # Make a TaskInstance that's deferred and waiting on it + # Make a TaskInstance that's deferred and waiting on it. A deferred task has + # already started running, so it has a start_date; set one so duration can be + # computed. Unlike set_state, handle_failure (used by the FAILED path) does not + # synthesize a missing start_date, matching the scheduler executor-event path. task_instance = create_task_instance( session=session, logical_date=timezone.utcnow(), state=State.DEFERRED ) task_instance.trigger_id = trigger.id + task_instance.start_date = now.subtract(seconds=10) session.commit() def get_xcoms(ti): @@ -335,6 +340,73 @@ def test_submit_event_task_end(mock_utcnow, session, create_task_instance, event assert actual_xcoms == expected_xcoms [email protected]( + ("retries", "expected_state", "expected_callback_type", "expect_history_row"), + [ + (1, TaskInstanceState.UP_FOR_RETRY, TaskInstanceState.UP_FOR_RETRY, True), + (0, TaskInstanceState.FAILED, TaskInstanceState.FAILED, False), + ], +) +@patch("airflow.callbacks.database_callback_sink.DatabaseCallbackSink.send") +def test_submit_event_task_end_failed_respects_retries( + mock_send, + session, + create_task_instance, + retries, + expected_state, + expected_callback_type, + expect_history_row, +): + """A trigger-emitted TaskFailedEvent should respect retry-eligibility: a deferred task with + retries remaining goes UP_FOR_RETRY (on_retry_callback), not straight to FAILED. + + On the retry path, the finished try must also be archived to task_instance_history so + prior-try log lookups keep working after the trigger ends the deferred try. + """ + trigger = Trigger(classpath="does.not.matter", kwargs={}) + session.add(trigger) + task_instance = create_task_instance( + session=session, + logical_date=timezone.utcnow(), + state=State.DEFERRED, + default_args={"retries": retries}, + ) + task_instance.trigger_id = trigger.id + task_instance.try_number = 1 + task_instance.max_tries = retries + old_ti_id = task_instance.id + session.commit() + + Trigger.submit_event(trigger.id, TaskFailedEvent(), session=session) + session.flush() + + ti = session.scalar(select(TaskInstance)) + assert ti.state == expected_state + + mock_send.assert_called_once() + request = mock_send.call_args.kwargs["callback"] + assert request.task_callback_type == expected_callback_type + + assert ti.next_method is None + assert ti.next_kwargs is None + assert ti.end_date is not None + + tih = session.scalars( + select(TaskInstanceHistory).where( + TaskInstanceHistory.dag_id == ti.dag_id, + TaskInstanceHistory.task_id == ti.task_id, + TaskInstanceHistory.run_id == ti.run_id, + ) + ).all() + if expect_history_row: + assert len(tih) == 1 + assert ti.id != old_ti_id + assert tih[0].task_instance_id == old_ti_id + else: + assert tih == [] + assert ti.id == old_ti_id + + @pytest.fixture def create_triggerer(): """Fixture factory which creates individual test Triggerer instances."""
