kaxil commented on code in PR #70685:
URL: https://github.com/apache/airflow/pull/70685#discussion_r3706954291
##########
airflow-core/src/airflow/models/trigger.py:
##########
@@ -509,33 +554,39 @@ 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 re-encoded with
the payload added)
+ is failed rather than resumed. That work happens in the scheduler and API
processes, which
+ handle every parked task instance in turn, so a single unusable payload
must not be able to
+ abort the caller.
+
: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.sdk.serde import serialize
from airflow.utils.state import TaskInstanceState
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
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)
-
- # 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
+ next_kwargs = _decode_next_kwargs(next_kwargs_raw)
+ # Add event to the plain dict, then serialize everything together so
nested
+ # non-primitive values get proper serde encoding.
+ next_kwargs["event"] = event.payload
+ # 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)
Review Comment:
One case this unit newly consumes: serde's `serialize` raises
`AttributeError` for any dict containing its reserved keys (`__classname__` /
`__schema_id__`), at any depth, and the HITL route stores `params_input`
verbatim with no key validation, so a response whose params contain such a key
blows up here. On main that request 500s, nothing is recorded, and the client
can fix and resubmit. With this change the response is recorded, the task is
routed to `__fail__`, the route returns 200 echoing the response, and the
`response_received` 409 guard blocks a corrected resubmission. A write-side 400
for serde-reserved keys in `params_input`, next to the existing option checks
in hitl.py, would keep that path recoverable (fine as a follow-up). Relatedly,
should this function report back whether it resumed or failed the task? The
HITL route currently has no way to know it just returned 200 for a response
that was effectively discarded.
##########
airflow-core/src/airflow/models/trigger.py:
##########
@@ -509,33 +554,39 @@ 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 re-encoded with
the payload added)
+ is failed rather than resumed. That work happens in the scheduler and API
processes, which
Review Comment:
This also runs in the triggerer, via `Trigger.submit_event`, and that caller
arguably had the worst pre-fix failure mode: `handle_events` leaves the event
unconfirmed on a raise, so the poison event is redelivered indefinitely. Maybe
"the triggerer, scheduler and API processes". On the same theme, the `:raise
ValueError:` note on `_decode_next_kwargs` documents only the isinstance
failure; serde and BaseSerialization decode errors propagate from it too, which
is exactly why this except has to stay broad. One sentence there would stop a
future reader narrowing it back down.
##########
airflow-core/src/airflow/models/trigger.py:
##########
@@ -500,6 +500,51 @@ 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.
+ """
+ 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,
*, session: Session) -> None:
+ """
+ Re-queue a task instance that cannot be resumed, so that a worker fails it.
+
+ Mirrors :meth:`Trigger.submit_failure`: the special ``__fail__``
next_method makes the worker
+ fail the task immediately, which runs its normal failure handling (retries
and callbacks
+ included). Leaving the task instance parked instead would strand it there,
as the event that
+ should have resumed it is already gone.
+ """
+ task_instance.next_method = TRIGGER_FAIL_REPR
+ task_instance.next_kwargs = {"error": reason}
Review Comment:
Worth also setting `"traceback": format_exception(type(exc), exc,
exc.__traceback__)` here like `submit_failure` does (means passing the
exception into this helper); `resume_execution` already renders that key into
the task log. As written the task log only shows the one-line reason while the
full traceback lands in the scheduler/triggerer/API log, which the DAG author
may not be able to see. It has to stay the list `format_exception` returns,
since the runtime does `"\n".join(traceback)`.
##########
airflow-core/src/airflow/models/trigger.py:
##########
@@ -509,33 +554,39 @@ 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 re-encoded with
the payload added)
+ is failed rather than resumed. That work happens in the scheduler and API
processes, which
+ handle every parked task instance in turn, so a single unusable payload
must not be able to
+ abort the caller.
+
: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.sdk.serde import serialize
from airflow.utils.state import TaskInstanceState
Review Comment:
This local import now shadows the module-level one from line 43
(`_fail_unresumable_task_instance` above uses the module-level name). Since the
line above is already touched, may as well drop it.
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -8077,6 +8077,60 @@ 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):
+ """
+ A parked task whose stored ``next_kwargs`` cannot be turned into a
dict is failed by the
+ sweep instead of aborting it, and the other timed-out tasks in the
same batch are still
+ resolved. The sweep runs on a recurring scheduler timer, so it has to
finish the batch.
+ """
+ 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())
+ self.job_runner.check_awaiting_input_timeouts(session=session)
Review Comment:
Related, though outside this diff: now that `handle_event_submit` can route
a TI to `__fail__` instead of resuming it, the sweep's response and defaults
branches still count that TI under `num_resolved`, so the summary log reports
it as "resolved (response/defaults)". Checking `ti.next_method ==
TRIGGER_FAIL_REPR` after the call and counting those separately would keep that
log line honest.
##########
airflow-core/src/airflow/models/trigger.py:
##########
@@ -509,33 +554,39 @@ 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 re-encoded with
the payload added)
+ is failed rather than resumed. That work happens in the scheduler and API
processes, which
+ handle every parked task instance in turn, so a single unusable payload
must not be able to
+ abort the caller.
+
: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.sdk.serde import serialize
from airflow.utils.state import TaskInstanceState
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
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)
-
- # 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
+ next_kwargs = _decode_next_kwargs(next_kwargs_raw)
+ # Add event to the plain dict, then serialize everything together so
nested
+ # non-primitive values get proper serde encoding.
+ next_kwargs["event"] = event.payload
+ # 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 process 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
processed "
Review Comment:
The try also spans `serialize` after the event payload is added, so a
failure caused by the incoming payload (a trigger yielding something serde
can't encode, or the reserved-keys case above) gets reported as the stored
next_kwargs being unreadable, pointing the author at DB state that was fine.
Splitting the decode from the inject-and-serialize step, with a
payload-specific message for the latter, would keep the diagnosis accurate.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
