Stop retries when the retry callback raises AirflowFailException

airflow-core/tests/unit/dag_processing/test_processor.py

@@ -2385,6 +2385,7 @@ def get_type_names(union_type):
             "GetHITLDetailResponse",
             "SetRenderedMapIndex",
             "SetExecutionTimeout",
+            "StartTaskFinalization",
             # AIP-103 task/asset store — Dag processor has no task execution context.
             "GetTaskStateStore",
             "SetTaskStateStore",

airflow-core/tests/unit/jobs/test_triggerer_job.py

@@ -3405,6 +3405,7 @@ def get_type_names(union_type):
             "CreateHITLDetailPayload",
             "SetRenderedMapIndex",
             "SetExecutionTimeout",
+            "StartTaskFinalization",
             "GetDag",
             # AIP-103 task store — triggerer has no task execution context.
             "GetTaskStateStore",

go-sdk/schema/supervisor-schema.json

@@ -4070,6 +4070,25 @@
       "title": "SkipDownstreamTasks",
       "type": "object"
     },
+    "StartTaskFinalization": {
+      "description": "Start bounded task finalization without reporting a task outcome.",
+      "properties": {
+        "retry": {
+          "$ref": "#/$defs/RetryTask"
+        },
+        "type": {
+          "const": "StartTaskFinalization",
+          "default": "StartTaskFinalization",
+          "title": "Type",
+          "type": "string"
+        }
+      },
+      "required": [
+        "retry"
+      ],
+      "title": "StartTaskFinalization",
+      "type": "object"
+    },
     "StartupDetails": {
       "properties": {
         "ti": {

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

@@ -1156,6 +1156,13 @@ class SetExecutionTimeout(BaseModel):
     type: Literal["SetExecutionTimeout"] = "SetExecutionTimeout"
 
 
+class StartTaskFinalization(BaseModel):
+    """Start bounded task finalization without reporting a task outcome."""
+
+    retry: RetryTask
+    type: Literal["StartTaskFinalization"] = "StartTaskFinalization"
+
+
 class TriggerDagRun(TriggerDAGRunPayload):
     dag_id: str
     run_id: Annotated[str, Field(title="Dag Run Id")]
@@ -1366,6 +1373,7 @@ class GetDag(BaseModel):
     | SetXCom
     | SkipDownstreamTasks
     | SetExecutionTimeout
+    | StartTaskFinalization
     | SucceedTask
     | ValidateInletsAndOutlets
     | TaskState

task-sdk/src/airflow/sdk/execution_time/schema/schema.json

@@ -4070,6 +4070,25 @@
       "title": "SkipDownstreamTasks",
       "type": "object"
     },
+    "StartTaskFinalization": {
+      "description": "Start bounded task finalization without reporting a task outcome.",
+      "properties": {
+        "retry": {
+          "$ref": "#/$defs/RetryTask"
+        },
+        "type": {
+          "const": "StartTaskFinalization",
+          "default": "StartTaskFinalization",
+          "title": "Type",
+          "type": "string"
+        }
+      },
+      "required": [
+        "retry"
+      ],
+      "title": "StartTaskFinalization",
+      "type": "object"
+    },
     "StartupDetails": {
       "properties": {
         "ti": {

task-sdk/src/airflow/sdk/execution_time/schema/versions/v2027_02_28.py

@@ -22,6 +22,8 @@
 
 - ``HITLDetailResponseResult``: the supervisor's reply when a running task reads its own
   Human-in-the-loop response.
+- ``StartTaskFinalization``: starts the supervisor's finalization timeout without choosing
+  a task outcome, allowing retry callbacks to decide whether the task should retry.
 """
 
 from __future__ import annotations

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

@@ -138,6 +138,7 @@
     SetTaskStateStore,
     SetXCom,
     SkipDownstreamTasks,
+    StartTaskFinalization,
     StartupDetails,
     SucceedTask,
     TaskBreadcrumbsResult,
@@ -1672,6 +1673,8 @@ class ActivitySubprocess(WatchedSubprocess):
         TaskState | SucceedTask | RetryTask | DeferTask | RescheduleTask | AwaitInputTask | None
     ) = attrs.field(default=None, init=False)
 
+    _retry_callback_fallback: RetryTask | None = attrs.field(default=None, init=False)
+
     _last_successful_heartbeat: float = attrs.field(default=0, init=False)
     _last_heartbeat_attempt: float = attrs.field(default=0, init=False)
 
@@ -1822,6 +1825,11 @@ def wait(self) -> int:
         return self._exit_code
 
     def update_task_state_if_needed(self):
+        if self._terminal_state is None and self._retry_callback_fallback is not None:
+            self._terminal_state = self._retry_callback_fallback.state
+            self._pending_terminal_state_msg = self._retry_callback_fallback
+            self._rendered_map_index = self._retry_callback_fallback.rendered_map_index
+        self._retry_callback_fallback = None
         if self._terminal_state != SERVER_TERMINATED and self._pending_terminal_state_msg is not None:
             if isinstance(self._pending_terminal_state_msg, (TaskState, RetryTask)):
                 self._send_terminal_state_msg(self._pending_terminal_state_msg)
@@ -2137,9 +2145,11 @@ def _handle_task_state(
         self, msg: TaskState | RetryTask, log: FilteringBoundLogger, req_id: int
     ) -> RequestResult:
         if self._terminal_state != SERVER_TERMINATED:
+            self._retry_callback_fallback = None
             self._terminal_state = msg.state
             self._pending_terminal_state_msg = msg
-            self._task_end_time_monotonic = time.monotonic()
+            if self._task_end_time_monotonic is None:
+                self._task_end_time_monotonic = time.monotonic()
             self._rendered_map_index = msg.rendered_map_index
         return None, {}
 
@@ -2177,6 +2187,16 @@ def _handle_set_execution_timeout(
         self._arm_execution_timeout()
         return None, {}
 
+    def _handle_start_task_finalization(
+        self, msg: StartTaskFinalization, log: FilteringBoundLogger, req_id: int
+    ) -> RequestResult:
+        if self._terminal_state != SERVER_TERMINATED:
+            self._retry_callback_fallback = msg.retry
+        if self._task_end_time_monotonic is None:
+            self._task_end_time_monotonic = time.monotonic()
+        self._execution_timeout_enforce_at = None
+        return None, {}
+
     def _arm_execution_timeout(self) -> None:
         # Zero means no timeout, as a zero timedelta does for a Python task.
         if not self._execution_timeout_seconds:
@@ -2486,6 +2506,7 @@ def _handle_clear_asset_state_store_by_uri(
                 register_request_method(SetTaskStateStore, _handle_set_task_state_store),
                 register_request_method(SkipDownstreamTasks, _handle_skip_downstream_tasks),
                 register_request_method(SetExecutionTimeout, _handle_set_execution_timeout),
+                register_request_method(StartTaskFinalization, _handle_start_task_finalization),
                 register_request_method(SucceedTask, _handle_finished_task),
                 register_request_method(TaskState, _handle_task_state),
                 register_request_method(TriggerDagRun, _handle_trigger_dag_run),
@@ -2699,7 +2720,11 @@ def start(  # type: ignore[override]
                 context = ti.get_template_context()
                 log = structlog.get_logger(logger_name="task")
 
-                state, msg, error = run(ti, context, log)
+                try:
+                    state, msg, error = run(ti, context, log)
+                finally:
+                    if supervisor._retry_callback_fallback is not None:
+                        supervisor.update_task_state_if_needed()
                 context["exception"] = error
                 try:
                     finalize(ti, state, context, log, error)

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

@@ -123,6 +123,7 @@
     SetRenderedFields,
     SetRenderedMapIndex,
     SkipDownstreamTasks,
+    StartTaskFinalization,
     StartupDetails,
     SucceedTask,
     TaskBreadcrumbsResult,
@@ -1804,6 +1805,24 @@ def run(
         # skipping the retry decision and every callback in ``finalize()``.
         msg, state, error = _handle_handler_failure(ti, e, log, context)
     finally:
+        if state == TaskInstanceState.UP_FOR_RETRY:
+            ti.state = state
+            context["exception"] = error
+            if ti.task.on_retry_callback and isinstance(msg, RetryTask):
+                try:
+                    SUPERVISOR_COMMS.send(StartTaskFinalization(retry=msg))
+                except Exception:
+                    log.exception("Failed to report start of task finalization to supervisor")
+            try:
+                _run_task_state_change_callbacks(ti.task, "on_retry_callback", context, log)
+            except AirflowFailException:
+                log.info("Retry callback raised AirflowFailException, failing the task instead of retrying")
+                state = TaskInstanceState.FAILED
+                msg = TaskState(
+                    state=TaskInstanceState.FAILED,
+                    end_date=ti.end_date,
+                    rendered_map_index=ti.rendered_map_index,
+                )
         # `state` may still be unset if an exception handler above raised before
         # binding it
         if state is not None:
@@ -2235,7 +2254,9 @@ def _run_task_state_change_callbacks(
     for i, callback in enumerate(getattr(task, kind)):
         try:
             create_executable_runner(callback, context_get_outlet_events(context), logger=log).run(context)
-        except Exception:
+        except Exception as error:
+            if kind == "on_retry_callback" and isinstance(error, AirflowFailException):
+                raise
             log.exception("Failed to run task callback", kind=kind, index=i, callback=callback)
 
 
@@ -2584,7 +2605,6 @@ 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)
         try:
             get_listener_manager().hook.on_task_instance_failed(
                 previous_state=TaskInstanceState.RUNNING, task_instance=ti, error=error