From a66b44686d7792e9f953b587c683ddf4f061cc66 Mon Sep 17 00:00:00 2001 From: rjgoyln Date: Thu, 20 Aug 2026 02:08:17 +0800 Subject: [PATCH] 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 --- airflow-core/newsfragments/71850.bugfix.rst | 1 + airflow-core/src/airflow/models/deadline.py | 7 ++ .../src/airflow/models/taskinstance.py | 16 +++- .../serialization/definitions/deadline.py | 5 +- .../src/airflow/serialization/encoders.py | 5 ++ .../tests/unit/models/test_deadline.py | 13 ++-- .../tests/unit/models/test_deadline_alert.py | 57 +++++++++++++- .../tests/unit/models/test_serialized_dag.py | 26 ++++++- .../tests/unit/models/test_taskinstance.py | 74 ++++++++++++++++++- .../src/airflow/sdk/definitions/deadline.py | 15 ++++ 10 files changed, 203 insertions(+), 16 deletions(-) create mode 100644 airflow-core/newsfragments/71850.bugfix.rst diff --git a/airflow-core/newsfragments/71850.bugfix.rst b/airflow-core/newsfragments/71850.bugfix.rst new file mode 100644 index 0000000000000..a8e5ebb341421 --- /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 77e3c682e7bf3..cc809659a8738 100644 --- a/airflow-core/src/airflow/models/deadline.py +++ b/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: diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index ff610f2146588..a42e5b8252e3d 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/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() diff --git a/airflow-core/src/airflow/serialization/definitions/deadline.py b/airflow-core/src/airflow/serialization/definitions/deadline.py index 968d69e3037fa..824d15814fe3a 100644 --- a/airflow-core/src/airflow/serialization/definitions/deadline.py +++ b/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: diff --git a/airflow-core/src/airflow/serialization/encoders.py b/airflow-core/src/airflow/serialization/encoders.py index b3aac16363d65..1ad61b6ce8804 100644 --- a/airflow-core/src/airflow/serialization/encoders.py +++ b/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 diff --git a/airflow-core/tests/unit/models/test_deadline.py b/airflow-core/tests/unit/models/test_deadline.py index 1771269b548a0..2eba97f6fc4e1 100644 --- a/airflow-core/tests/unit/models/test_deadline.py +++ b/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( diff --git a/airflow-core/tests/unit/models/test_deadline_alert.py b/airflow-core/tests/unit/models/test_deadline_alert.py index 56ed2f8cc865e..f957a602da97b 100644 --- a/airflow-core/tests/unit/models/test_deadline_alert.py +++ b/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) diff --git a/airflow-core/tests/unit/models/test_serialized_dag.py b/airflow-core/tests/unit/models/test_serialized_dag.py index f3d720ddcc611..c00f9095ddc38 100644 --- a/airflow-core/tests/unit/models/test_serialized_dag.py +++ b/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. diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 6a348a3f2ca26..bd7373491406d 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/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` diff --git a/task-sdk/src/airflow/sdk/definitions/deadline.py b/task-sdk/src/airflow/sdk/definitions/deadline.py index f3a14aeb9cc55..0bd83a94f8592 100644 --- a/task-sdk/src/airflow/sdk/definitions/deadline.py +++ b/task-sdk/src/airflow/sdk/definitions/deadline.py @@ -37,6 +37,13 @@ # 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): """ @@ -51,6 +58,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.""" @@ -89,6 +100,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): @@ -318,8 +331,10 @@ def register_custom_reference( # 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}; "