This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f0c13dc5f94 Fix task callbacks being skipped when
`TriggerDagRunOperator` gets a 404 (#70719)
f0c13dc5f94 is described below
commit f0c13dc5f948f7987728c571a43c0a7eeafbf521
Author: Kaxil Naik <[email protected]>
AuthorDate: Tue Aug 4 12:05:00 2026 +0100
Fix task callbacks being skipped when `TriggerDagRunOperator` gets a 404
(#70719)
* Fix task callbacks being skipped when TriggerDagRunOperator gets a 404
`run()` maps a task's outcome through a flat chain of `except` clauses, and
several of those clauses do real work: they call the API server, or
serialize
user-supplied values. Python does not offer an exception raised inside an
`except` clause to that clause's siblings, so when one of them raised, the
exception escaped `run()` entirely -- skipping the retry decision in
`_handle_current_task_failed()` and every callback, listener and failure
email
in `finalize()`.
The reported path: triggering a Dag that does not exist returns 404, which
`DagRunOperations.trigger` re-raises (it only special-cases the 409
already-exists case). The supervisor turns it into an `API_SERVER_ERROR`
response and `CommsDecoder._from_frame` raises `AirflowRuntimeError` -- from
inside `except DagRunTriggerException`, a few lines above the clause that
already handles `AirflowRuntimeError`. It is not limited to that path:
`_defer_task` and `_await_input_task` run `serde_serialize` over
user-supplied
kwargs, which raises `TypeError` for any value serde has no serializer for.
Split the function at the point where deciding the outcome ends and
reporting
it begins. `_run_task_and_map_outcome()` keeps the chain verbatim and
returns
the outcome; `run()` calls it, and its `except` now covers the handlers too.
The chain and the terminal-state `finally` block are untouched, so this is a
behaviour change rather than a reshuffle of existing lines.
`_handle_handler_failure()` keeps the chain's own classifications instead of
routing everything through the retry-count check, so `AirflowFailException`,
`AirflowSensorTimeout` and `AirflowTaskTerminated` still fail without
retrying
when they surface from a handler. It catches `Exception`, not
`BaseException`,
so `KeyboardInterrupt` still reaches `main()`'s exit-code-2 path -- the
supervisor's default termination signal is SIGINT, so swallowing it would
turn
an operator-initiated kill into an ordinary retry. If the failure path
itself
raises, it fails closed on a plain FAILED state rather than re-entering the
code that just failed and escaping again.
* Potential fix for pull request finding
Co-authored-by: Copilot Autofix powered by AI
<[email protected]>
---------
Co-authored-by: Copilot Autofix powered by AI
<[email protected]>
---
.../src/airflow/sdk/execution_time/task_runner.py | 98 ++++++-
.../task_sdk/execution_time/test_task_runner.py | 308 ++++++++++++++++++++-
2 files changed, 382 insertions(+), 24 deletions(-)
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 31d0831bbad..1ce848128d6 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -73,9 +73,12 @@ from airflow.sdk.definitions.mappedoperator import
MappedOperator
from airflow.sdk.definitions.param import process_params
from airflow.sdk.exceptions import (
AirflowException,
+ AirflowFailException,
AirflowInactiveAssetInInletOrOutletException,
AirflowRescheduleException,
AirflowRuntimeError,
+ AirflowSensorTimeout,
+ AirflowTaskTerminated,
AirflowTaskTimeout,
ErrorType,
TaskAwaitingInput,
@@ -238,6 +241,9 @@ class RuntimeTaskInstance(TaskInstance):
_terminal_state_send_failed: bool = False
"""True when the supervisor IPC send for a non-success terminal state
raised; signals main() to sys.exit(1) after finalize() so the supervisor
doesn't misclassify the run as SUCCESS via exit code 0."""
+ _failure_metrics_emitted: bool = False
+ """True once the failure counters have been recorded, so a second pass
through the failure path (one attempt raised part-way through) does not count
the same failure twice."""
+
_ti_context_from_server: Annotated[TIRunContext | None, Field(repr=False)]
= None
"""The Task Instance context from the API server, if any."""
@@ -1524,22 +1530,17 @@ def _await_input_task(
return msg, state
[email protected]_errors
-@detail_span("run")
-def run(
+def _run_task_and_map_outcome(
ti: RuntimeTaskInstance,
context: Context,
log: Logger,
) -> tuple[TaskInstanceState, ToSupervisor | None, BaseException | None]:
- """Run the task in this process."""
+ """Execute the task and map its outcome -- success or a handled exception
-- to a message and state."""
import signal
from airflow.sdk.exceptions import (
- AirflowFailException,
AirflowRescheduleException,
- AirflowSensorTimeout,
AirflowSkipException,
- AirflowTaskTerminated,
DagRunTriggerException,
DownstreamTasksSkipped,
TaskAwaitingInput,
@@ -1565,9 +1566,6 @@ def run(
state: TaskInstanceState | None = None
error: BaseException | None = None
- stats_tags = ti.stats_tags
- stats.incr("ti.start", tags=stats_tags)
-
try:
# First, clear the xcom data sent from server
if ti._ti_context_from_server and (keys_to_delete :=
ti._ti_context_from_server.xcom_keys_to_clear):
@@ -1691,6 +1689,32 @@ def run(
log.info("::group::Post Execute")
msg, state = _handle_current_task_failed(ti, e, log, context)
error = e
+
+ return state, msg, error
+
+
[email protected]_errors
+@detail_span("run")
+def run(
+ ti: RuntimeTaskInstance,
+ context: Context,
+ log: Logger,
+) -> tuple[TaskInstanceState, ToSupervisor | None, BaseException | None]:
+ """Run the task in this process."""
+ msg: ToSupervisor | None = None
+ state: TaskInstanceState | None = None
+ error: BaseException | None = None
+
+ stats_tags = ti.stats_tags
+ stats.incr("ti.start", tags=stats_tags)
+
+ try:
+ state, msg, error = _run_task_and_map_outcome(ti, context, log)
+ except Exception as e:
+ # Python does not offer an exception raised inside an ``except``
clause to that
+ # clause's siblings, so a handler that fails would otherwise escape
entirely --
+ # skipping the retry decision and every callback in ``finalize()``.
+ msg, state, error = _handle_handler_failure(ti, e, log, context)
finally:
# `state` may still be unset if an exception handler above raised
before
# binding it
@@ -1798,6 +1822,46 @@ def _evaluate_retry_policy(
return None
+def _handle_handler_failure(
+ ti: RuntimeTaskInstance, exception: Exception, log: Logger, context:
Context
+) -> tuple[RetryTask | TaskState, TaskInstanceState, Exception]:
+ """
+ Decide the outcome for an exception raised by one of the outcome handlers
themselves.
+
+ Handlers reach this by talking to the API server -- a missing Dag makes
+ ``_handle_trigger_dag_run`` raise ``AirflowRuntimeError`` -- or by
serializing
+ user-supplied values, since ``_defer_task`` and ``_await_input_task`` run
+ ``serde_serialize`` over kwargs and it raises ``TypeError`` for anything
it has no
+ serializer for.
+
+ The handler chain's own classifications are preserved rather than routing
everything
+ through the retry-count check, so an exception that means "do not retry"
still means
+ that when it surfaces from a handler.
+ """
+ log.exception("Task failed with exception")
+ if isinstance(exception, (AirflowFailException, AirflowSensorTimeout,
AirflowTaskTerminated)):
+ return _terminal_failure(ti), TaskInstanceState.FAILED, exception
+ try:
+ msg, state = _handle_current_task_failed(ti, exception, log, context)
+ except Exception:
+ # The failure path itself failed. Re-entering it would just fail again
and
+ # escape, losing the callbacks this whole function exists to preserve,
so fail
+ # closed on a plain FAILED state instead.
+ log.exception("Could not determine terminal state, failing closed")
+ return _terminal_failure(ti), TaskInstanceState.FAILED, exception
+ return msg, state, exception
+
+
+def _terminal_failure(ti: RuntimeTaskInstance) -> TaskState:
+ """Build a plain FAILED terminal message, bypassing any retry decision."""
+ ti.end_date = datetime.now(tz=timezone.utc)
+ return TaskState(
+ state=TaskInstanceState.FAILED,
+ end_date=ti.end_date,
+ rendered_map_index=ti.rendered_map_index,
+ )
+
+
def _handle_current_task_failed(
ti: RuntimeTaskInstance,
exception: BaseException,
@@ -1849,12 +1913,16 @@ def _finalize_task_failure(
end_date = datetime.now(tz=timezone.utc)
ti.end_date = end_date
- # Record operator and task instance failed metrics
- operator = ti.task.__class__.__name__
- stats_tags = ti.stats_tags
+ # Record operator and task instance failed metrics. One failure is one
increment even
+ # if this runs twice, which happens when a first pass raised after
counting -- see
+ # `_handle_handler_failure`.
+ if not ti._failure_metrics_emitted:
+ operator = ti.task.__class__.__name__
+ stats_tags = ti.stats_tags
- stats.incr("operator_failures", tags={**stats_tags, "operator_name":
operator})
- stats.incr("ti_failures", tags=stats_tags)
+ stats.incr("operator_failures", tags={**stats_tags, "operator_name":
operator})
+ stats.incr("ti_failures", tags=stats_tags)
+ ti._failure_metrics_emitted = True
if ti._ti_context_from_server and ti._ti_context_from_server.should_retry:
retry_kwargs: dict[str, Any] = {"end_date": end_date}
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 cbd194e35eb..cc9fb77e089 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
@@ -52,6 +52,7 @@ from airflow.api_fastapi.execution_api.routes.task_instances
import _emit_task_s
from airflow.listeners import hookimpl
from airflow.providers.standard.operators.python import PythonOperator
from airflow.providers.standard.operators.trigger_dagrun import
TriggerDagRunOperator
+from airflow.providers.standard.triggers.temporal import DateTimeTrigger
from airflow.sdk import (
DAG,
BaseOperator,
@@ -79,7 +80,13 @@ from airflow.sdk.bases.xcom import BaseXCom
from airflow.sdk.definitions._internal.types import NOTSET,
SET_DURING_EXECUTION, is_arg_set
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetUniqueKey,
AssetUriRef, Dataset, Model
from airflow.sdk.definitions.param import DagParam
-from airflow.sdk.definitions.retry_policy import ExceptionRetryPolicy,
RetryAction, RetryRule
+from airflow.sdk.definitions.retry_policy import (
+ ExceptionRetryPolicy,
+ RetryAction,
+ RetryDecision,
+ RetryPolicy,
+ RetryRule,
+)
from airflow.sdk.exceptions import (
AirflowException,
AirflowFailException,
@@ -133,6 +140,7 @@ from airflow.sdk.execution_time.comms import (
PreviousTIResult,
PrevSuccessfulDagRunResult,
RescheduleTask,
+ RetryTask,
SetAssetStateStoreByName,
SetAssetStateStoreByUri,
SetRenderedFields,
@@ -980,6 +988,214 @@ def test_defer_task_queue_assignment(
)
[email protected](
+ ("should_retry", "expected_state"),
+ [
+ (True, TaskInstanceState.UP_FOR_RETRY),
+ (False, TaskInstanceState.FAILED),
+ ],
+)
+def test_defer_with_unserializable_kwargs_honours_retries_and_callbacks(
+ should_retry, expected_state, create_runtime_ti, mock_supervisor_comms
+):
+ """
+ A task that defers with a non-serializable ``next_kwargs`` value must fail
like any
+ other task, rather than taking the whole run down.
+
+ ``_defer_task`` runs ``serde_serialize`` on the deferral kwargs, which
raises
+ ``TypeError`` for anything it has no serializer for (a file handle, a
client object,
+ a lambda). That raise happens inside ``run()``'s ``except TaskDeferred``
handler, so
+ before the fix it escaped ``run()`` without evaluating retries or running
callbacks,
+ skipping the retry decision and callbacks even without involving the API
server at
+ all.
+ """
+ callbacks_run = []
+
+ class _DeferWithBadKwargs(BaseOperator):
+ def execute(self, context):
+ raise TaskDeferred(
+ trigger=DateTimeTrigger(moment=timezone.datetime(2024, 11,
22)),
+ method_name="next",
+ # A live handle is the realistic version of this mistake.
+ kwargs={"client": object()},
+ )
+
+ task = _DeferWithBadKwargs(
+ task_id="defer_bad_kwargs",
+ on_failure_callback=lambda context: callbacks_run.append("failure"),
+ on_retry_callback=lambda context: callbacks_run.append("retry"),
+ )
+ ti = create_runtime_ti(
+ dag_id="test_defer_with_unserializable_kwargs",
+ run_id="test_run",
+ task=task,
+ should_retry=should_retry,
+ )
+
+ log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)
+ context = ti.get_template_context()
+
+ state, _, error = run(ti, context, log)
+
+ assert state == expected_state
+ assert isinstance(error, TypeError)
+
+ context["exception"] = error
+ finalize(ti, state, context, log, error)
+
+ assert callbacks_run == ["retry" if should_retry else "failure"]
+
+
+def test_handler_failure_keeps_non_retryable_exceptions_non_retryable(
+ create_runtime_ti, mock_supervisor_comms
+):
+ """
+ A handler that raises a non-retryable exception must still fail without
retrying.
+
+ `AirflowFailException` means "do not retry" wherever it is raised. Routing
handler
+ failures through `_handle_current_task_failed` would consult the retry
count instead
+ and hand back UP_FOR_RETRY, so `_handle_handler_failure` keeps the main
chain's
+ classification for these types.
+ """
+
+ class _FailingTrigger(DateTimeTrigger):
+ def serialize(self):
+ raise AirflowFailException("trigger cannot be serialized, do not
retry")
+
+ class _DeferWithFailingTrigger(BaseOperator):
+ def execute(self, context):
+ raise TaskDeferred(
+ trigger=_FailingTrigger(moment=timezone.datetime(2024, 11,
22)),
+ method_name="next",
+ )
+
+ task = _DeferWithFailingTrigger(task_id="defer_fail_exc")
+ # Retries are available; the exception type must still win.
+ ti = create_runtime_ti(
+ dag_id="test_handler_failure_non_retryable",
+ run_id="test_run",
+ task=task,
+ should_retry=True,
+ )
+
+ log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)
+
+ state, msg, error = run(ti, ti.get_template_context(), log)
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(msg, TaskState)
+ assert isinstance(error, AirflowFailException)
+
+
+def test_handler_failure_lets_keyboard_interrupt_propagate(create_runtime_ti,
mock_supervisor_comms):
+ """
+ A ``KeyboardInterrupt`` inside a handler must reach ``main()``, not become
a task failure.
+
+ The supervisor's default termination signal is SIGINT, so swallowing it
here would
+ convert an operator-initiated kill into an ordinary retry.
+ """
+
+ class _InterruptingTrigger(DateTimeTrigger):
+ def serialize(self):
+ raise KeyboardInterrupt
+
+ class _DeferWithInterrupt(BaseOperator):
+ def execute(self, context):
+ raise TaskDeferred(
+ trigger=_InterruptingTrigger(moment=timezone.datetime(2024,
11, 22)),
+ method_name="next",
+ )
+
+ task = _DeferWithInterrupt(task_id="defer_interrupt")
+ ti = create_runtime_ti(dag_id="test_handler_interrupt", run_id="test_run",
task=task)
+
+ log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)
+
+ with pytest.raises(KeyboardInterrupt):
+ run(ti, ti.get_template_context(), log)
+
+
+def
test_handler_failure_fails_closed_when_failure_path_also_fails(create_runtime_ti,
mock_supervisor_comms):
+ """
+ If the failure path itself raises, `run()` must still return a terminal
state.
+
+ A `retry_policy` is user code and `RetryDecision` is an unvalidated
dataclass, so a
+ policy handing back `retry_delay=30` (seconds, rather than a `timedelta`)
reaches
+ `_finalize_task_failure`, which calls `.total_seconds()` on it. That
`AttributeError`
+ is raised outside `_evaluate_retry_policy`'s own error handling, so it
propagates out
+ of the failure path. Re-entering that path with its own exception would
fail the same
+ way and escape `run()`, losing the callbacks this handling exists to
preserve.
+ """
+
+ class _BadDelayPolicy(RetryPolicy):
+ def evaluate(self, exception, try_number, max_tries, context=None):
+ # Seconds as an int, not a timedelta.
+ return RetryDecision(action=RetryAction.RETRY, retry_delay=30)
+
+ class _AlwaysFails(BaseOperator):
+ def execute(self, context):
+ raise RuntimeError("boom")
+
+ task = _AlwaysFails(task_id="bad_delay_policy",
retry_policy=_BadDelayPolicy())
+ ti = create_runtime_ti(
+ dag_id="test_handler_double_fault",
+ run_id="test_run",
+ task=task,
+ should_retry=True,
+ )
+
+ log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)
+
+ state, msg, error = run(ti, ti.get_template_context(), log)
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(msg, TaskState)
+ assert msg.state == TaskInstanceState.FAILED
+ # The terminal state must actually reach the supervisor, not merely be
returned.
+ assert any(call.kwargs.get("msg") is msg for call in
mock_supervisor_comms.send.call_args_list)
+ # `error` is the exception that actually ended the run, so a callback sees
the broken
+ # policy rather than a misleadingly clean task error. The task's own
failure stays
+ # reachable on the implicit exception chain.
+ assert isinstance(error, AttributeError)
+ assert isinstance(error.__context__, RuntimeError)
+
+
+def test_handler_failure_counts_the_failure_once(create_runtime_ti,
mock_supervisor_comms):
+ """
+ One failure is one increment, even when the failure path runs twice.
+
+ The first pass through `_finalize_task_failure` records the counters
before the broken
+ retry delay makes it raise, so counting again on the second pass would
report two
+ failures for a single task run.
+ """
+
+ class _BadDelayPolicy(RetryPolicy):
+ def evaluate(self, exception, try_number, max_tries, context=None):
+ return RetryDecision(action=RetryAction.RETRY, retry_delay=30)
+
+ class _AlwaysFails(BaseOperator):
+ def execute(self, context):
+ raise RuntimeError("boom")
+
+ task = _AlwaysFails(task_id="count_once", retry_policy=_BadDelayPolicy())
+ ti = create_runtime_ti(
+ dag_id="test_handler_failure_counts_once",
+ run_id="test_run",
+ task=task,
+ should_retry=True,
+ )
+
+ log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)
+ stats_backend = mock.MagicMock(spec=StatsLogger)
+
+ with mock.patch("airflow.sdk.execution_time.task_runner.stats",
stats_backend):
+ run(ti, ti.get_template_context(), log)
+
+ counted = [call.args[0] for call in stats_backend.incr.call_args_list if
call.args]
+ assert counted.count("ti_failures") == 1
+ assert counted.count("operator_failures") == 1
+
+
def test_run_downstream_skipped(mocked_parse, create_runtime_ti,
mock_supervisor_comms, listener_manager):
listener = TestTaskRunnerCallsListeners.CustomListener()
listener_manager(listener)
@@ -5153,10 +5369,10 @@ class TestTriggerDagRunOperator:
mock_supervisor_comms.assert_has_calls(expected_calls)
@time_machine.travel("2025-01-01 00:00:00", tick=False)
- def test_handle_trigger_dag_run_reraises_original_error(self,
create_runtime_ti, mock_supervisor_comms):
+ def test_handle_trigger_dag_run_surfaces_original_error(self,
create_runtime_ti, mock_supervisor_comms):
"""
- When an ``except`` handler in ``run()`` raises before binding
``state``,
- the original exception must propagate
+ When an ``except`` handler in ``run()`` raises, the original exception
must be
+ surfaced as the task failure, not an ``UnboundLocalError`` on
``state``.
"""
from airflow.providers.standard.operators.trigger_dagrun import
TriggerDagRunOperator
@@ -5169,7 +5385,7 @@ class TestTriggerDagRunOperator:
trigger_run_id="test_run_id",
)
ti = create_runtime_ti(
- dag_id="test_handle_trigger_dag_run_reraises_original_error",
+ dag_id="test_handle_trigger_dag_run_surfaces_original_error",
run_id="test_run",
task=task,
)
@@ -5183,11 +5399,85 @@ class TestTriggerDagRunOperator:
mock_supervisor_comms.send.side_effect = _send
- log = mock.MagicMock()
+ log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)
+
+ state, _, error = run(ti, ti.get_template_context(), log)
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(error, _TriggerSendError)
+
+ @pytest.mark.parametrize(
+ ("should_retry", "expected_state"),
+ [
+ (True, TaskInstanceState.UP_FOR_RETRY),
+ (False, TaskInstanceState.FAILED),
+ ],
+ )
+ @time_machine.travel("2025-01-01 00:00:00", tick=False)
+ def test_handle_trigger_dag_run_missing_dag_honours_retries_and_callbacks(
+ self, should_retry, expected_state, create_runtime_ti,
mock_supervisor_comms
+ ):
+ """
+ A 404 from the API server for a missing target Dag must be handled
like any other
+ task failure: the retry decision is made and ``finalize()`` fires the
callbacks.
+
+ Regression test for https://github.com/apache/airflow/issues/70683 --
the
+ ``AirflowRuntimeError`` was raised from inside ``run()``'s
+ ``except DagRunTriggerException`` handler, so it escaped ``run()``
without
+ evaluating retries or running ``on_failure_callback`` /
``on_retry_callback``.
+ """
+ callbacks_run = []
+
+ task = TriggerDagRunOperator(
+ task_id="test_task",
+ trigger_dag_id="this_dag_does_not_exist",
+ trigger_run_id="test_run_id",
+ on_failure_callback=lambda context:
callbacks_run.append("failure"),
+ on_retry_callback=lambda context: callbacks_run.append("retry"),
+ )
+ ti = create_runtime_ti(
+ dag_id="test_handle_trigger_dag_run_missing_dag",
+ run_id="test_run",
+ task=task,
+ should_retry=should_retry,
+ )
+
+ not_found = AirflowRuntimeError(
+ error=ErrorResponse(
+ error=ErrorType.API_SERVER_ERROR,
+ detail={
+ "status_code": 404,
+ "message": "Client error message: Dag with dag_id:
'this_dag_does_not_exist' not found",
+ },
+ )
+ )
+
+ def _send(msg=None, **kwargs):
+ if isinstance(msg, TriggerDagRun):
+ raise not_found
+ return mock.DEFAULT
+
+ mock_supervisor_comms.send.side_effect = _send
+
+ log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)
+ context = ti.get_template_context()
+
+ state, msg, error = run(ti, context, log)
+
+ assert state == expected_state
+ assert error is not_found
+ # The terminal message is what the server acts on, so assert the shape
too:
+ # RetryTask drives UP_FOR_RETRY, TaskState(FAILED) ends the try.
+ if should_retry:
+ assert isinstance(msg, RetryTask)
+ else:
+ assert isinstance(msg, TaskState)
+ assert msg.state == TaskInstanceState.FAILED
+
+ context["exception"] = error
+ finalize(ti, state, context, log, error)
- # The original error must surface, not UnboundLocalError on ``state``.
- with pytest.raises(_TriggerSendError):
- run(ti, ti.get_template_context(), log)
+ assert callbacks_run == ["retry" if should_retry else "failure"]
@pytest.mark.parametrize(
("allowed_states", "failed_states", "target_dr_state",
"expected_task_state"),
