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

Author: rjgoyln

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.

airflow-core/src/airflow/models/deadline.py

@@ -56,6 +56,11 @@
 
 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 classproperty:
     """
@@ -325,6 +330,7 @@ class BaseDeadlineReference(LoggingMixin, ABC):
 
         # Set of required kwargs - subclasses should override this.
         required_kwargs: set[str] = set()
+        evaluation_timing: str = DAGRUN_CREATED_TIMING
 
         @classproperty
         def reference_name(cls: Any) -> str:
@@ -411,6 +417,7 @@ class DagRunQueuedAtDeadline(BaseDeadlineReference):
         """A deadline that returns when a DagRun was queued."""
 
         required_kwargs = {"dag_id", "run_id"}
+        evaluation_timing: str = DAGRUN_QUEUED_TIMING
 
         @provide_session
         def _evaluate_with(self, *, session: Session, **kwargs: Any) -> datetime | None:

airflow-core/src/airflow/models/taskinstance.py

@@ -84,7 +84,12 @@
 from airflow.models.asset import AssetModel
 from airflow.models.base import Base, StringID, TaskInstanceDependencies
 from airflow.models.dag_version import DagVersion
-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
 
 # Import HITLDetail at runtime so SQLAlchemy can resolve the relationship
@@ -240,8 +245,13 @@ def _recalculate_dagrun_queued_at_deadlines(
         .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()
 

airflow-core/src/airflow/serialization/definitions/deadline.py

@@ -27,7 +27,7 @@
 from sqlalchemy import select
 
 from airflow._shared.timezones import timezone
-from airflow.models.deadline import classproperty
+from airflow.models.deadline import DAGRUN_CREATED_TIMING, DAGRUN_QUEUED_TIMING, classproperty
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import provide_session
 from airflow.utils.sqlalchemy import get_dialect_name
@@ -97,6 +97,7 @@ class SerializedBaseDeadlineReference(LoggingMixin, ABC):
         """Base class for all serialized Deadline implementations."""
 
         required_kwargs: set[str] = set()
+        evaluation_timing: str = DAGRUN_CREATED_TIMING
 
         @classproperty
         def reference_name(cls: Any) -> str:
@@ -180,6 +181,7 @@ class DagRunQueuedAtDeadline(SerializedBaseDeadlineReference):
         """A deadline that returns when a DagRun was queued."""
 
         required_kwargs = {"dag_id", "run_id"}
+        evaluation_timing: str = DAGRUN_QUEUED_TIMING
 
         @provide_session
         def _evaluate_with(self, *, session: Session, **kwargs: Any) -> datetime | None:
@@ -285,6 +287,7 @@ class SerializedCustomReference(SerializedBaseDeadlineReference):
 
         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:

airflow-core/src/airflow/serialization/encoders.py

@@ -26,6 +26,7 @@
 import pendulum
 
 from airflow._shared.module_loading import qualname
+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
@@ -277,6 +278,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

airflow-core/tests/unit/models/test_deadline.py

@@ -34,6 +34,8 @@
 from airflow.sdk import timezone
 from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback
 from airflow.sdk.definitions.deadline import (
+    DAGRUN_CREATED_TIMING,
+    DAGRUN_QUEUED_TIMING,
     AverageRuntimeDeadline,
     BaseDeadlineReference,
     DagRunLogicalDateDeadline,
@@ -769,14 +771,14 @@ def teardown_method(self):
         ],
     )
     @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
@@ -787,6 +789,7 @@ def test_register_custom_reference(self, timing, reference):
         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(

airflow-core/tests/unit/models/test_deadline_alert.py

@@ -24,10 +24,20 @@
 from sqlalchemy import select
 
 from airflow._shared.timezones import timezone
+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
@@ -39,6 +49,24 @@
 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()
 
@@ -236,11 +264,32 @@ def _evaluate_with(self, *, session, dag_id, run_id):
                 dag_id="test_dag",
             )
 
+    @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)
 

airflow-core/tests/unit/models/test_serialized_dag.py

@@ -41,7 +41,12 @@
 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
@@ -1265,6 +1270,25 @@ def test_write_dag_with_deadline_passes_schema_validation(self, testing_dag_bund
         # 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_does_not_mutate_caller_deadline_data(self, testing_dag_bundle, session):
         """write_dag must not rewrite the caller's LazyDeserializedDAG deadline in place.
 

airflow-core/tests/unit/models/test_taskinstance.py

@@ -91,14 +91,19 @@
 )
 from airflow.sdk.api.datamodels._generated import AssetEventResponse, AssetResponse
 from airflow.sdk.definitions.callback import AsyncCallback
-from airflow.sdk.definitions.deadline import DeadlineReference
+from airflow.sdk.definitions.deadline import (
+    DAGRUN_CREATED_TIMING,
+    DAGRUN_QUEUED_TIMING,
+    BaseDeadlineReference,
+    DeadlineReference,
+)
 from airflow.sdk.definitions.param import process_params
 from airflow.sdk.definitions.taskgroup import TaskGroup
 from airflow.sdk.execution_time.comms import AssetEventsResult
 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
@@ -4190,6 +4195,71 @@ def test_clear_task_instances_recalculates_dagrun_queued_deadlines(dag_maker, se
     assert recalculated_count == 2
 
 
+@pytest.mark.parametrize(
+    ("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, timing, should_recalculate
+):
+    """Test that clearing tasks recalculates custom deadlines registered as queued-anchored."""
+
+    class CustomReference(BaseDeadlineReference):
+        evaluation_timing = timing
+
+    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=encode_deadline_reference(CustomReference()),
+        interval=interval.total_seconds(),
+        callback_def={"path": f"{__name__}.empty_callback_for_deadline", "kwargs": {}},
+    )
+    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()
+    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 + interval
+    else:
+        assert deadline.deadline_time == original_deadline_time
+
+
 def test_get_dagrun_loaded_but_none_returns_dagrun(dag_maker, session):
     """
     Test that `get_dagrun()` fetches `DagRun` from DB when the `dag_run`