This is an automated email from the ASF dual-hosted git repository.
ferruzzi 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 646b92b8b9b Re-anchor custom queued deadline references when a Dag run
is cleared (#71850)
646b92b8b9b is described below
commit 646b92b8b9b66bd14b74382e868b3d8f06590f32
Author: rjgoyln <[email protected]>
AuthorDate: Thu Nov 12 03:18:56 2026 +0800
Re-anchor custom queued deadline references when a Dag run is cleared
(#71850)
* Re-anchor custom queued deadline references when a Dag run is cleared
The evaluation timing a reference is registered with only lived in the
in-memory registry of the process that ran the Dag file, so the scheduler
had no way to tell a queued-anchored reference from any other and fell
back to matching the built-in class name. Custom references registered as
DeadlineReference.TYPES.DAGRUN_QUEUED were silently left pointing at the
old queued_at.
closes: #71747
Co-authored-by: D. Ferruzzi <[email protected]>
---
airflow-core/newsfragments/71850.bugfix.rst | 1 +
airflow-core/src/airflow/models/deadline.py | 9 ++
airflow-core/src/airflow/models/serialized_dag.py | 44 +++---
airflow-core/src/airflow/models/taskinstance.py | 38 +++--
.../airflow/serialization/definitions/deadline.py | 12 +-
airflow-core/src/airflow/serialization/encoders.py | 5 +
airflow-core/tests/unit/models/test_deadline.py | 13 +-
.../tests/unit/models/test_deadline_alert.py | 57 ++++++-
.../tests/unit/models/test_serialized_dag.py | 64 +++++++-
.../tests/unit/models/test_taskinstance.py | 167 ++++++++++++++++++++-
task-sdk/src/airflow/sdk/definitions/deadline.py | 15 ++
11 files changed, 380 insertions(+), 45 deletions(-)
diff --git a/airflow-core/newsfragments/71850.bugfix.rst
b/airflow-core/newsfragments/71850.bugfix.rst
new file mode 100644
index 00000000000..a8e5ebb3414
--- /dev/null
+++ b/airflow-core/newsfragments/71850.bugfix.rst
@@ -0,0 +1 @@
+Custom deadline references registered as
``DeadlineReference.TYPES.DAGRUN_QUEUED`` are now re-anchored to the new queued
time when a Dag run is cleared, as the built-in
``DeadlineReference.DAGRUN_QUEUED_AT`` already was. Deadline alerts serialized
before this fix pick the behaviour up once their Dag is parsed again.
diff --git a/airflow-core/src/airflow/models/deadline.py
b/airflow-core/src/airflow/models/deadline.py
index b2c80d1b576..f7f0037c3a4 100644
--- a/airflow-core/src/airflow/models/deadline.py
+++ b/airflow-core/src/airflow/models/deadline.py
@@ -58,6 +58,11 @@ logger = logging.getLogger(__name__)
CALLBACK_METRICS_PREFIX = "deadline_alerts"
+# Must be in sync with the identically named constants in
airflow.sdk.definitions.deadline
+EVALUATION_TIMING_FIELD = "evaluation_timing"
+DAGRUN_CREATED_TIMING = "DAGRUN_CREATED"
+DAGRUN_QUEUED_TIMING = "DAGRUN_QUEUED"
+
class DeadlineDagRunProtocol(Protocol):
"""The Dag run interface deadline references are evaluated against."""
@@ -386,6 +391,8 @@ class ReferenceModels:
class BaseDeadlineReference(LoggingMixin, ABC):
"""Base class for all Deadline implementations."""
+ evaluation_timing: str = DAGRUN_CREATED_TIMING
+
@classproperty
def reference_name(cls: Any) -> str:
return cls.__name__
@@ -453,6 +460,8 @@ class ReferenceModels:
class DagRunQueuedAtDeadline(BaseDeadlineReference):
"""A deadline that returns when a DagRun was queued."""
+ evaluation_timing: str = DAGRUN_QUEUED_TIMING
+
def _evaluate_with(self, *, session: Session, dagrun:
DeadlineDagRunProtocol) -> datetime | None:
return dagrun.queued_at
diff --git a/airflow-core/src/airflow/models/serialized_dag.py
b/airflow-core/src/airflow/models/serialized_dag.py
index 6cab8defb81..63ca5a7b1e3 100644
--- a/airflow-core/src/airflow/models/serialized_dag.py
+++ b/airflow-core/src/airflow/models/serialized_dag.py
@@ -442,7 +442,7 @@ class SerializedDagModel(Base):
existing_deadline_uuids: list[str],
new_deadline_data: list[dict],
session: Session,
- ) -> tuple[dict[str, dict], dict[str, str | None]] | None:
+ ) -> tuple[dict[str, dict], dict[str, dict[str, Any]]] | None:
"""
Try to reuse existing deadline UUIDs if the deadline definitions
haven't changed.
@@ -451,11 +451,11 @@ class SerializedDagModel(Base):
:param existing_deadline_uuids: List of UUID strings from existing
serialized Dag
:param new_deadline_data: List of new deadline alert data dicts from
the Dag
:param session: Database session
- :return: Tuple of (uuid_mapping, name_updates) if all definitions
match, None if any
+ :return: Tuple of (uuid_mapping, row_updates) if all definitions
match, None if any
mismatch detected. ``uuid_mapping`` maps UUID string → new
deadline data dict.
- ``name_updates`` maps UUID string → new name **only** for entries
whose name
- changed relative to the existing DB row, so callers can issue
targeted UPDATEs
- and reliably detect whether any DB write occurred.
+ ``row_updates`` maps UUID string → the columns whose value
**changed** relative to
+ the existing DB row, so callers can issue targeted UPDATEs and
reliably detect
+ whether any DB write occurred.
"""
# defensive check for old 3.1.x format
if existing_deadline_uuids and not
isinstance(existing_deadline_uuids[0], str):
@@ -464,8 +464,12 @@ class SerializedDagModel(Base):
def _definitions_match(deadline_data: dict, existing:
DeadlineAlertModel) -> bool:
"""Check if raw deadline data matches an existing DeadlineAlert's
definition."""
+ # The reference is compared on the fields the stored row already
has: a field added to
+ # the serialized form by an upgrade would otherwise mark every
stored alert as changed,
+ # orphaning the deadlines of running Dag runs. Such a row is
refreshed in place below.
+ new_reference = deadline_data[DeadlineAlertFields.REFERENCE]
return (
- deadline_data[DeadlineAlertFields.REFERENCE] ==
existing.reference
+ all(new_reference.get(field) == value for field, value in
existing.reference.items())
and deadline_data[DeadlineAlertFields.INTERVAL] ==
existing.interval
and deadline_data[DeadlineAlertFields.CALLBACK] ==
existing.callback_def
)
@@ -483,7 +487,7 @@ class SerializedDagModel(Base):
matched_uuids: set[UUID] = set()
uuid_mapping: dict[str, dict] = {}
- name_updates: dict[str, str | None] = {}
+ row_updates: dict[str, dict[str, Any]] = {}
for deadline_alert in new_deadline_data:
deadline_data = deadline_alert.get(Encoding.VAR, deadline_alert)
@@ -499,9 +503,15 @@ class SerializedDagModel(Base):
uuid_mapping[uuid_str] = deadline_data
matched_uuids.add(existing_alert.id)
found_match = True
+ updates: dict[str, Any] = {}
new_name = deadline_data.get(DeadlineAlertFields.NAME)
if new_name != existing_alert.name:
- name_updates[uuid_str] = new_name
+ updates["name"] = new_name
+ new_reference =
deadline_data[DeadlineAlertFields.REFERENCE]
+ if new_reference != existing_alert.reference:
+ updates["reference"] = new_reference
+ if updates:
+ row_updates[uuid_str] = updates
break
if not found_match:
@@ -510,7 +520,7 @@ class SerializedDagModel(Base):
# to another deadline), so partial reuse would risk stale
cross-references.
return None
- return uuid_mapping, name_updates
+ return uuid_mapping, row_updates
@classmethod
def _create_deadline_alert_records(
@@ -653,7 +663,7 @@ class SerializedDagModel(Base):
serialized_dag_hash = _prefetched.dag_hash
dag_version = _prefetched.dag_version
- name_updated = False
+ deadline_rows_updated = False
reused_deadline_data: dict[str, dict] | None = None
if dag.data.get("dag", {}).get("deadline"):
# The deadline handling below rewrites data["dag"]["deadline"]
from a list of
@@ -679,17 +689,17 @@ class SerializedDagModel(Base):
)
if reuse_result is not None:
- deadline_uuid_mapping, name_updates = reuse_result
+ deadline_uuid_mapping, row_updates = reuse_result
# All deadlines matched — reuse the UUIDs to preserve hash.
- # Only issue UPDATE statements for rows whose name
actually changed to
+ # Only issue UPDATE statements for rows whose columns
actually changed to
# avoid unnecessary writes and to make the return value
accurate.
- for uuid_str, new_name in name_updates.items():
+ for uuid_str, values in row_updates.items():
session.execute(
update(DeadlineAlertModel)
.where(DeadlineAlertModel.id == UUID(uuid_str))
- .values(name=new_name)
+ .values(**values)
)
- name_updated = bool(name_updates)
+ deadline_rows_updated = bool(row_updates)
dag.data["dag"]["deadline"] = existing_deadline_uuids
reused_deadline_data = deadline_uuid_mapping
deadline_uuid_mapping = {}
@@ -720,8 +730,8 @@ class SerializedDagModel(Base):
dag_source_code=dag_source_code,
session=session,
)
- if name_updated or bundle_metadata_changed:
- # A write occurred — a deadline alert name update and/or a
bundle
+ if deadline_rows_updated or bundle_metadata_changed:
+ # A write occurred — a deadline alert row update and/or a
bundle
# metadata refresh — so report True so callers know the DB
changed.
return True
log.debug("Serialized DAG (%s) is unchanged. Skipping writing to
DB", dag.dag_id)
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index 3de44152302..c7773c4e416 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -101,7 +101,12 @@ from airflow.models.asset import AssetModel
from airflow.models.base import Base, StringID
from airflow.models.dag_version import DagVersion
from airflow.models.dagbag import DBDagBag
-from airflow.models.deadline import Deadline, ReferenceModels
+from airflow.models.deadline import (
+ DAGRUN_QUEUED_TIMING,
+ EVALUATION_TIMING_FIELD,
+ Deadline,
+ ReferenceModels,
+)
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
from airflow.models.dynamic_region import (
SENTINEL_REGION_ID,
@@ -284,8 +289,13 @@ def _recalculate_dagrun_queued_at_deadlines(dagrun:
DagRun, *, session: Session)
.where(
Deadline.dagrun_id == dagrun.id,
Deadline.missed == false(),
-
DeadlineAlertModel.reference[ReferenceModels.REFERENCE_TYPE_FIELD].as_string()
- == ReferenceModels.DagRunQueuedAtDeadline.__name__,
+ or_(
+
DeadlineAlertModel.reference[EVALUATION_TIMING_FIELD].as_string() ==
DAGRUN_QUEUED_TIMING,
+ # Alerts serialized before the timing was persisted only
identify the built-in
+ # queued reference by name.
+
DeadlineAlertModel.reference[ReferenceModels.REFERENCE_TYPE_FIELD].as_string()
+ == ReferenceModels.DagRunQueuedAtDeadline.__name__,
+ ),
)
).all()
@@ -299,9 +309,17 @@ def _recalculate_dagrun_queued_at_deadlines(dagrun:
DagRun, *, session: Session)
try:
decoded_alert = decode_deadline_alert_model(deadline_alert)
interval = resolve_deadline_alert_interval(decoded_alert,
session=session)
- except (ValueError, TypeError):
- # Either step can fail: the alert may hold a payload the decoder
refuses, or a
- # variable-backed interval may point at a Variable that is missing
or non-numeric.
+ new_deadline_time = decoded_alert.reference.evaluate_with(
+ session=session,
+ interval=interval,
+ dagrun=dagrun,
+ dag_id=dagrun.dag_id,
+ run_id=dagrun.run_id,
+ )
+ except Exception:
+ # Any step can fail: the alert may hold a payload the decoder
refuses, a
+ # variable-backed interval may point at a Variable that is missing
or non-numeric,
+ # and the evaluation is the Dag author's own code, which may raise
anything at all.
# Leave this deadline alone rather than failing the whole clear.
log.warning(
"Error while recalculating deadline %s for DagRun %s.%s,
leaving it unchanged",
@@ -312,14 +330,6 @@ def _recalculate_dagrun_queued_at_deadlines(dagrun:
DagRun, *, session: Session)
)
continue
- new_deadline_time = decoded_alert.reference.evaluate_with(
- session=session,
- interval=interval,
- dagrun=dagrun,
- dag_id=dagrun.dag_id,
- run_id=dagrun.run_id,
- )
-
if new_deadline_time is None:
continue
diff --git a/airflow-core/src/airflow/serialization/definitions/deadline.py
b/airflow-core/src/airflow/serialization/definitions/deadline.py
index c7f1f1ee111..3522e2f378f 100644
--- a/airflow-core/src/airflow/serialization/definitions/deadline.py
+++ b/airflow-core/src/airflow/serialization/definitions/deadline.py
@@ -27,7 +27,12 @@ import attrs
from sqlalchemy import select
from airflow._shared.timezones import timezone
-from airflow.models.deadline import _get_evaluation_kwargs, classproperty
+from airflow.models.deadline import (
+ DAGRUN_CREATED_TIMING,
+ DAGRUN_QUEUED_TIMING,
+ _get_evaluation_kwargs,
+ classproperty,
+)
from airflow.models.variable import Variable
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.session import provide_session
@@ -97,6 +102,8 @@ class SerializedReferenceModels:
class SerializedBaseDeadlineReference(LoggingMixin, ABC):
"""Base class for all serialized Deadline implementations."""
+ evaluation_timing: str = DAGRUN_CREATED_TIMING
+
@classproperty
def reference_name(cls: Any) -> str:
return cls.__name__
@@ -161,6 +168,8 @@ class SerializedReferenceModels:
class DagRunQueuedAtDeadline(SerializedBaseDeadlineReference):
"""A deadline that returns when a DagRun was queued."""
+ evaluation_timing: str = DAGRUN_QUEUED_TIMING
+
def _evaluate_with(self, *, session: Session, dagrun:
DeadlineDagRunProtocol) -> datetime | None:
return dagrun.queued_at
@@ -261,6 +270,7 @@ class SerializedReferenceModels:
def __init__(self, inner_ref):
self.inner_ref = inner_ref
+ self.evaluation_timing = getattr(inner_ref, "evaluation_timing",
DAGRUN_CREATED_TIMING)
@property
def reference_name(self) -> str:
diff --git a/airflow-core/src/airflow/serialization/encoders.py
b/airflow-core/src/airflow/serialization/encoders.py
index 2100cedb3b8..7d3d93cdd67 100644
--- a/airflow-core/src/airflow/serialization/encoders.py
+++ b/airflow-core/src/airflow/serialization/encoders.py
@@ -27,6 +27,7 @@ import pendulum
from airflow._shared.module_loading import qualname
from airflow._shared.serialization import CLASSNAME, DATA, VERSION
+from airflow.models.deadline import DAGRUN_CREATED_TIMING,
EVALUATION_TIMING_FIELD
from airflow.partition_mappers.base import PartitionMapper as
CorePartitionMapper
from airflow.partition_mappers.wait_policy import WaitPolicy as CoreWaitPolicy
from airflow.partition_mappers.window import Window as CoreWindow
@@ -349,6 +350,10 @@ def encode_deadline_reference(ref) -> dict[str, Any]:
serialized = ref.serialize_reference()
+ # Added here rather than in serialize_reference() because custom
references are expected to
+ # override that method, and an override would drop the timing.
+ serialized[EVALUATION_TIMING_FIELD] = getattr(ref, "evaluation_timing",
DAGRUN_CREATED_TIMING)
+
# Custom types (not built-in) need __class_path so the decoder can look
them up.
# Unlike built-in types which are looked up in SerializedReferenceModels,
# custom types are resolved at deserialization time from the classes
registered
diff --git a/airflow-core/tests/unit/models/test_deadline.py
b/airflow-core/tests/unit/models/test_deadline.py
index 3848a89c7c9..29d7efb9ff0 100644
--- a/airflow-core/tests/unit/models/test_deadline.py
+++ b/airflow-core/tests/unit/models/test_deadline.py
@@ -36,6 +36,8 @@ from airflow.providers.standard.operators.empty import
EmptyOperator
from airflow.sdk import serde, timezone
from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback
from airflow.sdk.definitions.deadline import (
+ DAGRUN_CREATED_TIMING,
+ DAGRUN_QUEUED_TIMING,
AverageRuntimeDeadline,
BaseDeadlineReference,
DagRunLogicalDateDeadline,
@@ -790,14 +792,14 @@ class TestCustomDeadlineReference:
],
)
@pytest.mark.parametrize(
- "timing",
+ ("timing", "expected_evaluation_timing"),
[
- pytest.param(None, id="default_timing"),
- pytest.param(DeadlineReference.TYPES.DAGRUN_CREATED,
id="dagrun_created"),
- pytest.param(DeadlineReference.TYPES.DAGRUN_QUEUED,
id="dagrun_queued"),
+ pytest.param(None, DAGRUN_CREATED_TIMING, id="default_timing"),
+ pytest.param(DeadlineReference.TYPES.DAGRUN_CREATED,
DAGRUN_CREATED_TIMING, id="dagrun_created"),
+ pytest.param(DeadlineReference.TYPES.DAGRUN_QUEUED,
DAGRUN_QUEUED_TIMING, id="dagrun_queued"),
],
)
- def test_register_custom_reference(self, timing, reference):
+ def test_register_custom_reference(self, timing,
expected_evaluation_timing, reference):
if timing is None:
result = DeadlineReference.register_custom_reference(reference)
expected_timing = DeadlineReference.TYPES.DAGRUN_CREATED
@@ -808,6 +810,7 @@ class TestCustomDeadlineReference:
assert result is reference
assert hasattr(DeadlineReference, reference.__name__)
assert getattr(DeadlineReference, reference.__name__).__class__ is
reference
+ assert reference.evaluation_timing == expected_evaluation_timing
assert_correct_timing(reference, expected_timing)
assert_builtin_types_unchanged(
diff --git a/airflow-core/tests/unit/models/test_deadline_alert.py
b/airflow-core/tests/unit/models/test_deadline_alert.py
index 54a138b64a0..90b85586d2d 100644
--- a/airflow-core/tests/unit/models/test_deadline_alert.py
+++ b/airflow-core/tests/unit/models/test_deadline_alert.py
@@ -26,10 +26,20 @@ from sqlalchemy import select
from airflow._shared.timezones import timezone
from airflow.exceptions import RemovedInAirflow4Warning
+from airflow.models import deadline as deadline_models
+from airflow.models.deadline import ReferenceModels
from airflow.models.deadline_alert import DeadlineAlert
from airflow.models.serialized_dag import SerializedDagModel
-from airflow.sdk.definitions.deadline import BaseDeadlineReference,
DeadlineReference
+from airflow.sdk.definitions.deadline import (
+ DAGRUN_CREATED_TIMING,
+ DAGRUN_QUEUED_TIMING,
+ EVALUATION_TIMING_FIELD,
+ REFERENCE_TYPE_FIELD,
+ BaseDeadlineReference,
+ DeadlineReference,
+)
from airflow.serialization.definitions.deadline import
SerializedReferenceModels
+from airflow.serialization.encoders import encode_deadline_reference
from tests_common.test_utils import db
from unit.models import DEFAULT_DATE
@@ -41,6 +51,24 @@ DEADLINE_CALLBACK = {"path": "test.callback"}
SERIALIZED_DAG_ID = "serialized_dag_uuid"
+class CustomQueuedReference(BaseDeadlineReference):
+ """A queued-anchored custom reference which overrides serialize_reference,
as the docs suggest."""
+
+ evaluation_timing = DAGRUN_QUEUED_TIMING
+
+ def serialize_reference(self) -> dict:
+ return {REFERENCE_TYPE_FIELD: self.reference_name}
+
+
+def test_core_timing_constants_match_the_sdk():
+ """Core cannot import these from the SDK, so guard the copies against
drift."""
+ assert (
+ deadline_models.EVALUATION_TIMING_FIELD,
+ deadline_models.DAGRUN_CREATED_TIMING,
+ deadline_models.DAGRUN_QUEUED_TIMING,
+ ) == (EVALUATION_TIMING_FIELD, DAGRUN_CREATED_TIMING, DAGRUN_QUEUED_TIMING)
+
+
def _clean_db():
db.clear_db_deadline_alert()
@@ -298,11 +326,32 @@ class TestDeadlineAlert:
"run_id": "test_run",
}
+ @pytest.mark.parametrize(
+ ("reference", "expected_timing"),
+ [
+ pytest.param(DeadlineReference.DAGRUN_QUEUED_AT,
DAGRUN_QUEUED_TIMING, id="sdk_queued"),
+ pytest.param(DeadlineReference.DAGRUN_LOGICAL_DATE,
DAGRUN_CREATED_TIMING, id="sdk_created"),
+ pytest.param(CustomQueuedReference(), DAGRUN_QUEUED_TIMING,
id="custom_queued"),
+ pytest.param(
+
SerializedReferenceModels.SerializedCustomReference(CustomQueuedReference()),
+ DAGRUN_QUEUED_TIMING,
+ id="wrapped_custom_queued",
+ ),
+ pytest.param(
+ SerializedReferenceModels.DagRunQueuedAtDeadline(),
+ DAGRUN_QUEUED_TIMING,
+ id="serialized_queued",
+ ),
+ pytest.param(
+ ReferenceModels.DagRunQueuedAtDeadline(),
DAGRUN_QUEUED_TIMING, id="core_legacy_queued"
+ ),
+ ],
+ )
+ def test_encoded_reference_carries_evaluation_timing(self, reference,
expected_timing):
+ assert encode_deadline_reference(reference)[EVALUATION_TIMING_FIELD]
== expected_timing
+
def test_core_deadline_reference_treated_as_builtins(self):
"""Test that refs from airflow.models.deadline are still treated as
builtins."""
- from airflow.models.deadline import ReferenceModels
- from airflow.serialization.encoders import encode_deadline_reference
-
ref = ReferenceModels.DagRunLogicalDateDeadline()
serialized = encode_deadline_reference(ref)
diff --git a/airflow-core/tests/unit/models/test_serialized_dag.py
b/airflow-core/tests/unit/models/test_serialized_dag.py
index 7b47f70fdf2..7bec095c890 100644
--- a/airflow-core/tests/unit/models/test_serialized_dag.py
+++ b/airflow-core/tests/unit/models/test_serialized_dag.py
@@ -42,7 +42,12 @@ from airflow.providers.standard.operators.empty import
EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator
from airflow.sdk import DAG, Asset, AssetAlias, task as task_decorator
from airflow.sdk.definitions.callback import AsyncCallback
-from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineReference
+from airflow.sdk.definitions.deadline import (
+ DAGRUN_QUEUED_TIMING,
+ EVALUATION_TIMING_FIELD,
+ DeadlineAlert,
+ DeadlineReference,
+)
from airflow.serialization.dag_dependency import DagDependency
from airflow.serialization.definitions.dag import SerializedDAG
from airflow.serialization.serialized_objects import DagSerialization,
LazyDeserializedDAG
@@ -1371,6 +1376,63 @@ class TestSerializedDagModel:
# Must not raise: the stored UUID-reference form has to satisfy the
serialized Dag schema.
DagSerialization.validate_schema(result.data)
+ def test_write_dag_persists_evaluation_timing(self, testing_dag_bundle,
session):
+ """The stored alert must record when the reference is anchored, not
just its class name."""
+ dag_id = "test_deadline_evaluation_timing"
+ dag = DAG(
+ dag_id=dag_id,
+ deadline=DeadlineAlert(
+ reference=DeadlineReference.DAGRUN_QUEUED_AT,
+ interval=timedelta(minutes=5),
+ callback=AsyncCallback(empty_callback_for_deadline),
+ ),
+ )
+ EmptyOperator(task_id="task1", dag=dag)
+ sync_dag_to_db(dag, session=session)
+ session.commit()
+
+ serialized_dag = session.scalar(select(SDM).where(SDM.dag_id ==
dag_id))
+ alert = session.scalar(select(DAM).where(DAM.serialized_dag_id ==
serialized_dag.id))
+ assert alert.reference[EVALUATION_TIMING_FIELD] == DAGRUN_QUEUED_TIMING
+
+ def test_write_dag_refreshes_alert_predating_a_serialized_reference_field(
+ self, testing_dag_bundle, session
+ ):
+ """An alert stored before a reference field existed keeps its row and
gains the field."""
+ dag_id = "test_deadline_reference_field_added"
+ dag = DAG(
+ dag_id=dag_id,
+ deadline=DeadlineAlert(
+ reference=DeadlineReference.DAGRUN_QUEUED_AT,
+ interval=timedelta(minutes=5),
+ callback=AsyncCallback(empty_callback_for_deadline),
+ ),
+ )
+ EmptyOperator(task_id="task1", dag=dag)
+ sync_dag_to_db(dag, session=session)
+ session.commit()
+
+ def _count_versions() -> int:
+ return session.scalar(
+
select(func.count()).select_from(DagVersion).where(DagVersion.dag_id == dag_id)
+ )
+
+ alert = session.scalar(select(DAM))
+ alert_id = alert.id
+ versions_before = _count_versions()
+
+ # Mimic a row written by a version that did not serialize the
evaluation timing yet.
+ alert.reference = {k: v for k, v in alert.reference.items() if k !=
EVALUATION_TIMING_FIELD}
+ session.commit()
+
+ sync_dag_to_db(dag, session=session)
+ session.commit()
+
+ alerts = session.scalars(select(DAM)).all()
+ assert [stored.id for stored in alerts] == [alert_id]
+ assert alerts[0].reference[EVALUATION_TIMING_FIELD] ==
DAGRUN_QUEUED_TIMING
+ assert _count_versions() == versions_before
+
def test_write_dag_does_not_mutate_caller_deadline_data(self,
testing_dag_bundle, session):
"""write_dag must not rewrite the caller's LazyDeserializedDAG
deadline in place.
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index 04cbc99da90..da91ad8d18f 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -40,7 +40,7 @@ from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session, load_only
from sqlalchemy.orm.attributes import set_committed_value
-from airflow import settings
+from airflow import plugins_manager, settings
from airflow._shared.observability.metrics.base_stats_logger import StatsLogger
from airflow._shared.observability.traces import new_dagrun_trace_carrier,
new_task_run_carrier
from airflow._shared.timezones import timezone
@@ -96,7 +96,13 @@ from airflow.sdk import (
from airflow.sdk.api.datamodels._generated import AssetEventResponse,
AssetResponse
from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.definitions.callback import AsyncCallback
-from airflow.sdk.definitions.deadline import DeadlineReference,
VariableInterval
+from airflow.sdk.definitions.deadline import (
+ DAGRUN_CREATED_TIMING,
+ DAGRUN_QUEUED_TIMING,
+ BaseDeadlineReference,
+ DeadlineReference,
+ VariableInterval,
+)
from airflow.sdk.definitions.param import process_params
from airflow.sdk.definitions.taskgroup import TaskGroup
from airflow.sdk.execution_time.comms import AssetEventsResult
@@ -104,7 +110,7 @@ from airflow.sdk.serde import serialize
from airflow.serialization.definitions.assets import SerializedAsset
from airflow.serialization.definitions.baseoperator import
SerializedBaseOperator
from airflow.serialization.definitions.dag import SerializedDAG
-from airflow.serialization.encoders import ensure_serialized_asset
+from airflow.serialization.encoders import encode_deadline_reference,
ensure_serialized_asset
from airflow.serialization.serialized_objects import OperatorSerialization,
create_scheduler_operator
from airflow.ti_deps.dep_context import DepContext
from airflow.ti_deps.dependencies_deps import REQUEUEABLE_DEPS, RUNNING_DEPS
@@ -5034,6 +5040,161 @@ def
test_clear_task_instances_skips_deadline_with_unresolvable_interval(dag_make
assert args[0] == deadline.id
[email protected](
+ ("timing", "should_recalculate"),
+ [
+ pytest.param(DAGRUN_QUEUED_TIMING, True, id="queued"),
+ pytest.param(DAGRUN_CREATED_TIMING, False, id="created"),
+ ],
+)
+def test_clear_task_instances_recalculates_custom_queued_deadlines(
+ dag_maker, session, monkeypatch, timing, should_recalculate
+):
+ """Test that clearing tasks recalculates custom deadlines registered as
queued-anchored."""
+ # An offset of its own, so the new deadline can only come from the
reference's own evaluation
+ # and not from the queued_at the built-in reference would return.
+ offset = datetime.timedelta(minutes=30)
+
+ class CustomReference(BaseDeadlineReference):
+ evaluation_timing = timing
+
+ def _evaluate_with(self, *, session, dagrun):
+ return dagrun.queued_at + offset
+
+ reference = encode_deadline_reference(CustomReference())
+ # Decoding the alert to resolve its interval goes through the plugin
registry, which is where
+ # a custom reference has to be registered for the scheduler to resolve it
at all.
+ monkeypatch.setattr(
+ plugins_manager,
+ "get_deadline_references_plugins",
+ lambda: {reference["__class_path"]: CustomReference},
+ )
+
+ with dag_maker(
+ dag_id="test_recalculate_custom_deadlines",
+ schedule=datetime.timedelta(days=1),
+ ) as dag:
+ EmptyOperator(task_id="task_1")
+
+ dag_run = dag_maker.create_dagrun()
+ ti = dag_run.get_task_instance("task_1", session=session)
+ ti.set_state(TaskInstanceState.SUCCESS, session=session)
+
+ original_queued_at = timezone.utcnow() - datetime.timedelta(hours=2)
+ dag_run.queued_at = original_queued_at
+ session.flush()
+
+ interval = datetime.timedelta(hours=1)
+ deadline_alert = DeadlineAlertModel(
+ serialized_dag_id=session.scalar(
+ select(SerializedDagModel.id).where(SerializedDagModel.dag_id ==
dag.dag_id)
+ ),
+ reference=reference,
+ interval=serialize(interval),
+ callback_def=serialize(AsyncCallback(empty_callback_for_deadline)),
+ )
+ session.add(deadline_alert)
+ session.flush()
+
+ original_deadline_time = original_queued_at + offset + interval
+ session.add(
+ Deadline(
+ dagrun_id=dag_run.id,
+ deadline_alert_id=deadline_alert.id,
+ deadline_time=original_deadline_time,
+ callback=AsyncCallback(empty_callback_for_deadline),
+ dag_id=dag_run.dag_id,
+ )
+ )
+ session.flush()
+
+ tis = session.scalars(select(TI).where(TI.dag_id == dag.dag_id, TI.run_id
== dag_run.run_id)).all()
+ clear_task_instances(tis, session)
+
+ dag_run = session.scalar(select(DagRun).where(DagRun.id == dag_run.id))
+ deadline = session.scalar(select(Deadline).where(Deadline.dagrun_id ==
dag_run.id))
+
+ if should_recalculate:
+ assert deadline.deadline_time == dag_run.queued_at + offset + interval
+ else:
+ assert deadline.deadline_time == original_deadline_time
+
+
+def test_clear_task_instances_skips_deadline_whose_reference_raises(dag_maker,
session, monkeypatch):
+ """A clear must survive a custom reference that raises while it is being
re-anchored.
+
+ The evaluation is the Dag author's own code and is reached straight from
the clear API, so
+ anything it raises would otherwise surface there as an HTTP 500.
+ """
+
+ class RaisingReference(BaseDeadlineReference):
+ evaluation_timing = DAGRUN_QUEUED_TIMING
+
+ def _evaluate_with(self, *, session, dagrun):
+ raise RuntimeError("deliberate failure inside a custom deadline
reference")
+
+ reference = encode_deadline_reference(RaisingReference())
+ monkeypatch.setattr(
+ plugins_manager,
+ "get_deadline_references_plugins",
+ lambda: {reference["__class_path"]: RaisingReference},
+ )
+
+ with dag_maker(
+ dag_id="test_recalculate_raising_reference",
+ schedule=datetime.timedelta(days=1),
+ ) as dag:
+ EmptyOperator(task_id="task_1")
+
+ dag_run = dag_maker.create_dagrun()
+ ti = dag_run.get_task_instance("task_1", session=session)
+ ti.set_state(TaskInstanceState.SUCCESS, session=session)
+
+ original_queued_at = timezone.utcnow() - datetime.timedelta(hours=2)
+ dag_run.queued_at = original_queued_at
+ session.flush()
+
+ interval = datetime.timedelta(hours=1)
+ deadline_alert = DeadlineAlertModel(
+ serialized_dag_id=session.scalar(
+ select(SerializedDagModel.id).where(SerializedDagModel.dag_id ==
dag.dag_id)
+ ),
+ reference=reference,
+ interval=serialize(interval),
+ callback_def=serialize(AsyncCallback(empty_callback_for_deadline)),
+ )
+ session.add(deadline_alert)
+ session.flush()
+
+ original_deadline_time = original_queued_at + interval
+ session.add(
+ Deadline(
+ dagrun_id=dag_run.id,
+ deadline_alert_id=deadline_alert.id,
+ deadline_time=original_deadline_time,
+ callback=AsyncCallback(empty_callback_for_deadline),
+ dag_id=dag_run.dag_id,
+ )
+ )
+ session.flush()
+
+ tis = session.scalars(select(TI).where(TI.dag_id == dag.dag_id, TI.run_id
== dag_run.run_id)).all()
+
+ with mock.patch("airflow.models.taskinstance.log") as mock_log:
+ clear_task_instances(tis, session)
+
+ dag_run = session.scalar(select(DagRun).where(DagRun.id == dag_run.id))
+ assert dag_run.queued_at > original_queued_at
+
+ deadline = session.scalar(select(Deadline).where(Deadline.dagrun_id ==
dag_run.id))
+ assert deadline.deadline_time == original_deadline_time
+
+ assert mock_log.warning.call_count == 1
+ msg, *args = mock_log.warning.call_args.args
+ assert "Error while recalculating deadline" in msg
+ assert args[0] == deadline.id
+
+
def test_get_dagrun_loaded_but_none_returns_dagrun(dag_maker, session):
"""
Test that `get_dagrun()` fetches `DagRun` from DB when the `dag_run`
diff --git a/task-sdk/src/airflow/sdk/definitions/deadline.py
b/task-sdk/src/airflow/sdk/definitions/deadline.py
index 61dbe09b57b..b0915acc5b3 100644
--- a/task-sdk/src/airflow/sdk/definitions/deadline.py
+++ b/task-sdk/src/airflow/sdk/definitions/deadline.py
@@ -38,6 +38,13 @@ logger = logging.getLogger(__name__)
# Field name used in serialization - must be in sync with
SerializedReferenceModels.REFERENCE_TYPE_FIELD
REFERENCE_TYPE_FIELD = "reference_type"
+# Field name holding the evaluation timing of a reference in its serialized
form.
+EVALUATION_TIMING_FIELD = "evaluation_timing"
+
+# Values stored under EVALUATION_TIMING_FIELD; they mirror the
DeadlineReference.TYPES names.
+DAGRUN_CREATED_TIMING = "DAGRUN_CREATED"
+DAGRUN_QUEUED_TIMING = "DAGRUN_QUEUED"
+
class BaseDeadlineReference(ABC):
"""
@@ -53,6 +60,10 @@ class BaseDeadlineReference(ABC):
attribute of an ``AirflowPlugin``; see
:external:doc:`howto/deadline-alerts`.
"""
+ # Set by ``register_custom_reference`` and persisted by the serializer,
because the in-memory
+ # ``DeadlineReference.TYPES`` registry only exists in the process that ran
the Dag file.
+ evaluation_timing: str = DAGRUN_CREATED_TIMING
+
@property
def reference_name(self) -> str:
"""Return the class name as the reference identifier."""
@@ -91,6 +102,8 @@ class DagRunLogicalDateDeadline(BaseDeadlineReference):
class DagRunQueuedAtDeadline(BaseDeadlineReference):
"""A deadline that returns when a DagRun was queued."""
+ evaluation_timing: str = DAGRUN_QUEUED_TIMING
+
@dataclass
class FixedDatetimeDeadline(BaseDeadlineReference):
@@ -344,8 +357,10 @@ class DeadlineReference:
# Add to appropriate deadline_reference_type classification
if deadline_reference_type is cls.TYPES.DAGRUN_CREATED:
cls.TYPES.DAGRUN_CREATED = cls.TYPES.DAGRUN_CREATED +
(reference_class,)
+ reference_class.evaluation_timing = DAGRUN_CREATED_TIMING
elif deadline_reference_type is cls.TYPES.DAGRUN_QUEUED:
cls.TYPES.DAGRUN_QUEUED = cls.TYPES.DAGRUN_QUEUED +
(reference_class,)
+ reference_class.evaluation_timing = DAGRUN_QUEUED_TIMING
else:
raise ValueError(
f"Invalid deadline reference type {deadline_reference_type}; "
