From 332d11f78ff03f9a57256a5f126729f7c31077e1 Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Wed, 5 Aug 2026 20:47:57 +0800 Subject: [PATCH] Fail task instances whose stored next_kwargs cannot be processed (#70685) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fail task instances whose stored next_kwargs cannot be processed handle_event_submit decoded a task instance's stored next_kwargs and assumed the result was a dict. Neither assumption held: the decode caught only four exception types, so anything the BaseSerialization fallback raised escaped, and the isinstance check sat under TYPE_CHECKING, so it never ran at runtime. Both escape as exceptions from a function whose callers walk every waiting task instance in one pass — the scheduler's timeout sweep and two API routes — so a single unusable payload aborted the whole batch. Decode through a helper that checks its result, and guard decode, event insertion and re-encode together. A task instance whose payload cannot be processed is re-queued to fail through the existing __fail__ path, so its normal retry and callback handling still runs, instead of being left parked for the next sweep to trip over again. Generated-by: Claude Opus 5 (1M context) following the guidelines at https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions * Report why a task instance could not be resumed, and to whom The Dag author sees only the task log, so the traceback has to travel in next_kwargs the way submit_failure already sends it; the process log where this runs is often not theirs to read. Decode and re-encode also fail for different reasons: blaming the stored kwargs for a payload the trigger just yielded points the author at database state that was never at fault. The sweep's summary counted an unresumable task as resolved. * Update airflow-core/src/airflow/models/trigger.py Co-authored-by: Amogh Desai * Fix ruff D205/D213 on _fail_unresumable_task_instance docstring One-line summary on the second line + blank line before the description, so the docstring satisfies both D205 (blank between summary and description) and D213 (summary on second line). Static checks were failing on ruff for this. --------- Co-authored-by: Amogh Desai Co-authored-by: Rahul Vats (cherry picked from commit 2b7a0be359cd7bd9da9054aeac6762b14e286762) --- .../src/airflow/jobs/scheduler_job_runner.py | 21 +++- airflow-core/src/airflow/models/trigger.py | 112 +++++++++++++++--- .../tests/unit/jobs/test_scheduler_job.py | 52 ++++++++ .../tests/unit/models/test_trigger.py | 67 +++++++++++ 4 files changed, 233 insertions(+), 19 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 9114ff979b918..fb73845501947 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -3414,8 +3414,10 @@ def check_awaiting_input_timeouts( num_resolved = 0 num_failed = 0 + num_unresumable = 0 for ti in timed_out_tis: hitl_detail = ti.hitl_detail + resuming = True if hitl_detail is not None and hitl_detail.responded_at is not None: # A response landed just before the deadline; resume with it. handle_event_submit( @@ -3423,7 +3425,6 @@ def check_awaiting_input_timeouts( task_instance=ti, session=session, ) - num_resolved += 1 elif hitl_detail is not None and hitl_detail.defaults is not None: # Apply the configured defaults as the response, then resume to success. hitl_detail.chosen_options = list(hitl_detail.defaults) @@ -3439,7 +3440,6 @@ def check_awaiting_input_timeouts( task_instance=ti, session=session, ) - num_resolved += 1 else: # No defaults and no response: resume into execute_complete with a timeout # failure event so the operator raises HITLTimeoutError (matching the old @@ -3455,16 +3455,29 @@ def check_awaiting_input_timeouts( task_instance=ti, session=session, ) + resuming = False + + # ``handle_event_submit`` routes a task instance it could not process to + # ``__fail__`` instead of resuming it. That is neither of the outcomes the + # branches above intended, so it is counted on its own rather than being + # reported as resolved. + if ti.next_method == TRIGGER_FAIL_REPR: + num_unresumable += 1 + elif resuming: + num_resolved += 1 + else: num_failed += 1 # Flush within the retry block so both branches persist consistently (the defaults # branch already flushes via handle_event_submit; the fail branch relies on this). session.flush() - if num_resolved or num_failed: + if num_resolved or num_failed or num_unresumable: self.log.info( - "AWAITING_INPUT timeout sweep: %i resolved (response/defaults), %i failed", + "AWAITING_INPUT timeout sweep: %i resolved (response/defaults), %i failed, " + "%i could not be resumed", num_resolved, num_failed, + num_unresumable, ) # [START find_and_purge_task_instances_without_heartbeats] diff --git a/airflow-core/src/airflow/models/trigger.py b/airflow-core/src/airflow/models/trigger.py index 9d8c195e9c924..c7039750b4eae 100644 --- a/airflow-core/src/airflow/models/trigger.py +++ b/airflow-core/src/airflow/models/trigger.py @@ -500,6 +500,55 @@ def get_sorted_triggers( return result +def _decode_next_kwargs(next_kwargs_raw: Any) -> dict[str, Any]: + """ + Decode the stored ``next_kwargs`` of a task instance into a plain dict. + + Deserialize with serde first to provide a compat layer if there are mixed serialized + (BaseSerialisation and serde) data, which can happen if a deferred task resumes after upgrade. + + The result is checked here rather than assumed, so callers never have to trust the shape of + what comes back out of the stored payload. + + :raise ValueError: The payload did not decode to a dict. + :raise Exception: Whatever the two decoders raise on a payload they cannot read -- the stored + blob is arbitrary, so the set is open and callers have to treat it as such. + """ + from airflow.sdk.serde import deserialize + + try: + next_kwargs = deserialize(next_kwargs_raw) + except (ImportError, KeyError, AttributeError, TypeError): + from airflow.serialization.serialized_objects import BaseSerialization + + next_kwargs = BaseSerialization.deserialize(next_kwargs_raw) + + if not isinstance(next_kwargs, dict): + raise ValueError(f"next_kwargs decoded to {type(next_kwargs).__name__}, expected a dict") + return next_kwargs + + +def _fail_unresumable_task_instance( + task_instance: TaskInstance, reason: str, exc: BaseException, *, session: Session +) -> None: + """ + Route through ``__fail__`` so a worker fails the task normally instead of stranding it. + + Mirrors ``Trigger.submit_failure``: without this the task is left with no event to resume it. + Traceback goes into ``next_kwargs`` as a list -- the only channel reaching the task log -- + since ``format_exception`` returns it that way and the runtime joins it. + """ + task_instance.next_method = TRIGGER_FAIL_REPR + task_instance.next_kwargs = { + "error": reason, + "traceback": format_exception(type(exc), exc, exc.__traceback__), + } + task_instance.trigger_id = None + task_instance.state = TaskInstanceState.SCHEDULED + task_instance.scheduled_dttm = timezone.utcnow() + session.flush() + + @singledispatch def handle_event_submit(event: TriggerEvent, *, task_instance: TaskInstance, session: Session) -> None: """ @@ -509,33 +558,66 @@ def handle_event_submit(event: TriggerEvent, *, task_instance: TaskInstance, ses as well as its state to scheduled. It also adds the event's payload into the kwargs for the task. + A task instance whose stored kwargs cannot be decoded, or which the event payload cannot be + encoded into, is failed rather than resumed. This runs in the triggerer, the scheduler and the + API processes, each of which handles every waiting task instance in one pass, so a single + unusable payload must not be able to abort the caller. The triggerer had the worst of it: an + event whose submit raised was left unconfirmed and redelivered indefinitely. + + Failing the task instance is not free for every caller: a Human-in-the-loop response whose + ``params_input`` serde cannot encode is now recorded and discarded rather than rejected, which + the submitter cannot retry. That wants validating on the write side; tracked at + https://github.com/apache/airflow/issues/71036 + :param task_instance: The task instance to handle the submit event for. :param session: The session to be used for the database callback sink. """ - from airflow.sdk.serde import deserialize, serialize - from airflow.utils.state import TaskInstanceState + from airflow.sdk.serde import serialize next_kwargs_raw = task_instance.next_kwargs or {} - # deserialize first to provide a compat layer if there are mixed serialized (BaseSerialisation and serde) data - # which can happen if a deferred task resumes after upgrade + # Decoding and re-encoding fail for different reasons and are reported separately: blaming the + # stored kwargs for a payload the trigger just yielded would point the author at DB state that + # was never the problem. try: - next_kwargs = deserialize(next_kwargs_raw) - except (ImportError, KeyError, AttributeError, TypeError): - from airflow.serialization.serialized_objects import BaseSerialization - - next_kwargs = BaseSerialization.deserialize(next_kwargs_raw) + next_kwargs = _decode_next_kwargs(next_kwargs_raw) + except Exception as exc: + log.exception( + "Could not decode the stored next_kwargs of %s; failing it instead of resuming it", + task_instance, + ) + _fail_unresumable_task_instance( + task_instance, + "Could not resume the task: its stored next_kwargs could not be decoded " + f"({type(exc).__name__}: {exc})", + exc, + session=session, + ) + return # Add event to the plain dict, then serialize everything together so nested # non-primitive values get proper serde encoding. - if TYPE_CHECKING: - assert isinstance(next_kwargs, dict) next_kwargs["event"] = event.payload + try: + # Re-serialize using serde. The Execution API version converter + # (ModifyDeferredTaskKwargsToJsonValue) handles converting this to + # BaseSerialization format when serving old workers. + serialized_next_kwargs = serialize(next_kwargs) + except Exception as exc: + log.exception( + "Could not serialize the event payload for %s; failing it instead of resuming it", + task_instance, + ) + _fail_unresumable_task_instance( + task_instance, + f"Could not resume the task: the event payload could not be serialized " + f"({type(exc).__name__}: {exc})", + exc, + session=session, + ) + return - # Re-serialize using serde. The Execution API version converter - # (ModifyDeferredTaskKwargsToJsonValue) handles converting this to - # BaseSerialization format when serving old workers. - task_instance.next_kwargs = serialize(next_kwargs) + task_instance.next_kwargs = serialized_next_kwargs # Remove ourselves as its trigger task_instance.trigger_id = None diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 439869d33c7b0..804befe0ebae2 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -7787,6 +7787,58 @@ def test_awaiting_input_timeout_without_defaults_fails(self, dag_maker): assert ti.next_method == "execute_complete" assert ti.next_kwargs["event"]["error_type"] == "timeout" + def test_awaiting_input_timeout_sweep_survives_unusable_next_kwargs(self, dag_maker): + """The sweep finishes its batch even when one task's stored kwargs cannot be read.""" + session = settings.Session() + with dag_maker( + dag_id="test_awaiting_input_bad_kwargs", + start_date=DEFAULT_DATE, + schedule="@once", + session=session, + ): + EmptyOperator(task_id="dummy1") + dr_bad = dag_maker.create_dagrun() + dr_good = dag_maker.create_dagrun( + run_id="good", logical_date=DEFAULT_DATE + datetime.timedelta(seconds=1) + ) + ti_bad = dr_bad.get_task_instance("dummy1", session=session) + ti_good = dr_good.get_task_instance("dummy1", session=session) + for ti in (ti_bad, ti_good): + ti.state = State.AWAITING_INPUT + ti.trigger_timeout = timezone.utcnow() - datetime.timedelta(seconds=60) + ti.next_method = "execute_complete" + ti.next_kwargs = {} + # Stored kwargs that no longer decode into a dict (here: a class name that is not allowed + # for deserialization, which the legacy fallback cannot read either). + ti_bad.next_kwargs = {"__classname__": "not.allowed.Thing", "__version__": 1, "__data__": {}} + session.add( + HITLDetail( + ti_id=ti_good.id, + options=["Approve", "Reject"], + subject="approve?", + defaults=["Approve"], + multiple=False, + params={}, + ) + ) + session.flush() + + self.job_runner = SchedulerJobRunner(job=Job()) + mock_log = mock.MagicMock(spec=logging.Logger) + with mock.patch.object(SchedulerJobRunner, "log", mock_log): + self.job_runner.check_awaiting_input_timeouts(session=session) + + session.refresh(ti_bad) + session.refresh(ti_good) + assert ti_bad.state == State.SCHEDULED + assert ti_bad.next_method == "__fail__" + assert "error" in ti_bad.next_kwargs + assert ti_good.state == State.SCHEDULED + assert ti_good.next_method == "execute_complete" + assert ti_good.next_kwargs["event"]["chosen_options"] == ["Approve"] + # The one it could not resume is reported as such rather than counted as resolved. + assert mock_log.info.call_args.args[1:] == (1, 0, 1) + def test_retry_on_db_error_when_update_timeout_triggers(self, dag_maker, testing_dag_bundle, session): """ Tests that it will retry on DB error like deadlock when updating timeout triggers. diff --git a/airflow-core/tests/unit/models/test_trigger.py b/airflow-core/tests/unit/models/test_trigger.py index 2785c48c1dd36..797f092218970 100644 --- a/airflow-core/tests/unit/models/test_trigger.py +++ b/airflow-core/tests/unit/models/test_trigger.py @@ -35,6 +35,7 @@ from airflow.models.asset import AssetEvent, AssetModel, AssetWatcherModel from airflow.models.callback import Callback, TriggererCallback from airflow.models.taskinstancehistory import TaskInstanceHistory +from airflow.models.trigger import handle_event_submit from airflow.models.xcom import XComModel from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk.definitions.callback import AsyncCallback @@ -257,6 +258,72 @@ def test_submit_event_no_n_plus_one_for_assets(_, session, asset_count, expected Trigger.submit_event(trigger_id, TriggerEvent("payload"), session=session) +@pytest.mark.parametrize( + "stored_next_kwargs", + [ + # Decoding blows up: serde rejects the class name, and the BaseSerialization fallback then + # trips over the missing legacy keys. + pytest.param( + {"__classname__": "not.allowed.Thing", "__version__": 1, "__data__": {}}, + id="undecodable", + ), + # Decodes cleanly, but not into a dict: legacy encoding of a bare datetime. + pytest.param({"__type": "datetime", "__var": 1735689600.0}, id="not-a-dict"), + ], +) +def test_handle_event_submit_fails_task_with_unusable_next_kwargs( + session, create_task_instance, stored_next_kwargs +): + """ + Tests that stored kwargs which cannot be turned into a dict fail the task instance instead of + raising out of ``handle_event_submit``. Its callers walk every waiting task instance in one + pass, so one unusable payload must not abort them. + """ + task_instance = create_task_instance( + session=session, logical_date=timezone.utcnow(), state=State.DEFERRED + ) + task_instance.next_method = "execute_complete" + task_instance.next_kwargs = stored_next_kwargs + session.flush() + + handle_event_submit(TriggerEvent("payload"), task_instance=task_instance, session=session) + + session.refresh(task_instance) + assert task_instance.state == State.SCHEDULED + assert task_instance.next_method == "__fail__" + assert task_instance.trigger_id is None + assert "event" not in task_instance.next_kwargs + assert "stored next_kwargs could not be decoded" in task_instance.next_kwargs["error"] + # The traceback reaches the task log only through next_kwargs, and the runtime joins the list. + assert isinstance(task_instance.next_kwargs["traceback"], list) + + +def test_handle_event_submit_fails_task_when_the_event_payload_cannot_be_serialized( + session, create_task_instance +): + """A payload serde cannot encode is reported as such, not as unreadable stored kwargs. + + Blaming the stored kwargs would point the Dag author at database state that was never the + problem. + """ + task_instance = create_task_instance( + session=session, logical_date=timezone.utcnow(), state=State.DEFERRED + ) + task_instance.next_method = "execute_complete" + task_instance.next_kwargs = {} + session.flush() + + # serde refuses any dict carrying its reserved keys, at any depth. + handle_event_submit( + TriggerEvent({"__classname__": "anything"}), task_instance=task_instance, session=session + ) + + session.refresh(task_instance) + assert task_instance.state == State.SCHEDULED + assert task_instance.next_method == "__fail__" + assert "event payload could not be serialized" in task_instance.next_kwargs["error"] + + def test_submit_failure(session, create_task_instance): """ Tests that failures submitted to a trigger fail their dependent