diff --git a/task-sdk/newsfragments/60172.bugfix.rst b/task-sdk/newsfragments/60172.bugfix.rst new file mode 100644 index 0000000000000..a466856cdf579 --- /dev/null +++ b/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. diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 6e0b0766be380..7abe915673dd7 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/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) diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 88424bfe53e62..b5b754e71b109 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/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"""