Fix AirflowFailException in on_retry_callback not preventing retries

task-sdk/newsfragments/60172.bugfix.rst

@@ -0,0 +1 @@
+``AirflowFailException`` raised inside ``on_retry_callback`` now correctly prevents task retries and marks the task as failed.

task-sdk/src/airflow/sdk/execution_time/task_runner.py

@@ -1355,7 +1355,7 @@ def _on_term(signum, frame):
         # Same metric with tagging
         Stats.incr("ti.finish", tags={**stats_tags, "state": state.value})
 
-        if msg:
+        if msg and state != TaskInstanceState.UP_FOR_RETRY:
             SUPERVISOR_COMMS.send(msg=msg)
 
     # Return the message to make unit tests easier too
@@ -1538,10 +1538,14 @@ def _run_task_state_change_callbacks(
     context: Context,
     log: Logger,
 ) -> None:
+    from airflow.sdk.exceptions import AirflowFailException
+
     callback: Callable[[Context], None]
     for i, callback in enumerate(getattr(task, kind)):
         try:
             create_executable_runner(callback, context_get_outlet_events(context), logger=log).run(context)
+        except AirflowFailException:
+            raise
         except Exception:
             log.exception("Failed to run task callback", kind=kind, index=i, callback=callback)
 
@@ -1747,6 +1751,7 @@ def finalize(
     context: Context,
     log: Logger,
     error: BaseException | None = None,
+    msg: ToSupervisor | None = None,
 ):
     # Record task duration metrics for all terminal states
     if ti.start_date and ti.end_date:
@@ -1801,15 +1806,38 @@ def finalize(
         except Exception:
             log.exception("error calling listener")
     elif state == TaskInstanceState.UP_FOR_RETRY:
-        _run_task_state_change_callbacks(task, "on_retry_callback", context, log)
+        from airflow.sdk.exceptions import AirflowFailException
+
+        try:
+            _run_task_state_change_callbacks(task, "on_retry_callback", context, log)
+        except AirflowFailException as e:
+            # If on_retry_callback raises AirflowFailException, the task should not retry.
+            log.info("on_retry_callback raised AirflowFailException; marking task as failed without retry")
+            state = TaskInstanceState.FAILED
+            ti.state = state
+            error = e
+            SUPERVISOR_COMMS.send(
+                msg=TaskState(
+                    state=TaskInstanceState.FAILED,
+                    end_date=ti.end_date,
+                    rendered_map_index=ti.rendered_map_index,
+                )
+            )
+            _run_task_state_change_callbacks(task, "on_failure_callback", context, log)
+            if error and task.email_on_failure and task.email:
+                _send_error_email_notification(task, ti, context, error, log)
+        else:
+            # No AirflowFailException: proceed with retry as originally planned.
+            if msg:
+                SUPERVISOR_COMMS.send(msg=msg)
+            if error and task.email_on_retry and task.email:
+                _send_error_email_notification(task, ti, context, error, log)
         try:
             get_listener_manager().hook.on_task_instance_failed(
                 previous_state=TaskInstanceState.RUNNING, task_instance=ti, error=error
             )
         except Exception:
             log.exception("error calling listener")
-        if error and task.email_on_retry and task.email:
-            _send_error_email_notification(task, ti, context, error, log)
     elif state == TaskInstanceState.FAILED:
         _run_task_state_change_callbacks(task, "on_failure_callback", context, log)
         try:
@@ -1871,9 +1899,9 @@ def main():
                 bundle_name=ti.bundle_instance.name,
                 bundle_version=ti.bundle_instance.version,
             ):
-                state, _, error = run(ti, context, log)
+                state, run_msg, error = run(ti, context, log)
                 context["exception"] = error
-                finalize(ti, state, context, log, error)
+                finalize(ti, state, context, log, error, msg=run_msg)
         except KeyboardInterrupt:
             log.exception("Ctrl-c hit")
             sys.exit(2)

task-sdk/tests/task_sdk/execution_time/test_task_runner.py

@@ -4360,6 +4360,53 @@ class CustomOperator(BaseOperator):
             expected_exception_logs.insert(index, calls)
         assert log.exception.mock_calls == expected_exception_logs
 
+    def test_airflow_fail_exception_in_on_retry_callback_marks_task_as_failed(
+        self, create_runtime_ti, mock_supervisor_comms
+    ):
+        """AirflowFailException raised in on_retry_callback must prevent retries and mark task as FAILED."""
+        from airflow.sdk.exceptions import AirflowFailException
+
+        failure_callback_called = []
+        retry_callback_called = []
+
+        def retry_callback(context):
+            retry_callback_called.append(True)
+            raise AirflowFailException("No more retries!")
+
+        def failure_callback(context):
+            failure_callback_called.append(True)
+
+        class FailingOperator(BaseOperator):
+            def execute(self, context):
+                raise AirflowException("Task failed")
+
+        task = FailingOperator(
+            task_id="task",
+            on_retry_callback=retry_callback,
+            on_failure_callback=failure_callback,
+        )
+        runtime_ti = create_runtime_ti(dag_id="dag", task=task, should_retry=True)
+        log = mock.MagicMock()
+        context = runtime_ti.get_template_context()
+        state, run_msg, error = run(runtime_ti, context, log)
+        finalize(runtime_ti, state, context, log, error, msg=run_msg)
+
+        assert runtime_ti.state == TaskInstanceState.FAILED
+        assert retry_callback_called == [True]
+        assert failure_callback_called == [True]
+
+        # The supervisor must receive TaskState(FAILED), not RetryTask
+        sent_msgs = [
+            call.args[0] if call.args else call.kwargs.get("msg")
+            for call in mock_supervisor_comms.send.call_args_list
+        ]
+        from airflow.sdk.execution_time.comms import RetryTask, TaskState
+
+        assert not any(isinstance(m, RetryTask) for m in sent_msgs), (
+            "RetryTask must not be sent when on_retry_callback raises AirflowFailException"
+        )
+        assert any(isinstance(m, TaskState) and m.state == TaskInstanceState.FAILED for m in sent_msgs)
+
 
 class TestTriggerDagRunOperator:
     """Tests to verify various aspects of TriggerDagRunOperator"""