Copilot commented on code in PR #72559:
URL: https://github.com/apache/airflow/pull/72559#discussion_r3940367843
##########
airflow-core/src/airflow/models/deadline.py:
##########
@@ -42,18 +42,110 @@
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.session import provide_session
from airflow.utils.sqlalchemy import UtcDateTime, get_dialect_name
-from airflow.utils.state import CallbackState
+from airflow.utils.state import CallbackState, TaskInstanceState
if TYPE_CHECKING:
from sqlalchemy.orm import Session
from sqlalchemy.sql import ColumnElement
from airflow.models.callback import CallbackDefinitionProtocol
from airflow.models.deadline_alert import DeadlineAlert
+ from airflow.models.taskinstance import TaskInstance
logger = logging.getLogger(__name__)
+
+def get_task_deadline_alerts(task: Any) -> list[Any] | None:
+ """
+ Return the deadline alerts declared on a task, whether stored on the
operator or in partial kwargs.
+
+ A task's deadline definition travels with the operator. For a mapped task
the definition is kept
+ in ``partial_kwargs`` (the ``SerializedMappedOperator`` has no
``deadline`` attribute), so both
+ locations must be consulted. A single alert is normalized to a one-element
list.
+
+ :param task: Serialized operator (or partial kwargs holder) to inspect
+ :return: List of deadline alerts, or None when the task declares none
+ """
+ deadline_alerts = getattr(task, "deadline", None)
+ if not deadline_alerts:
+ partial_kwargs = getattr(task, "partial_kwargs", None) or {}
+ deadline_alerts = partial_kwargs.get("deadline")
+ if deadline_alerts is None:
+ return None
+ if isinstance(deadline_alerts, list):
+ return deadline_alerts
+ return [deadline_alerts]
+
+
+def create_deadlines_for_task_instance(
+ *,
+ deadline_alerts: list[Any],
+ task_instance: TaskInstance,
+ bundle_name: str | None,
+ session: Session,
+) -> int:
+ """
+ Create Deadline rows for each of a task instance's DeadlineAlerts.
+
+ Used whenever a runnable task instance comes into existence: at DagRun
creation for regular
+ tasks, and at expansion time for each ``map_index`` of a mapped task.
+
+ :param deadline_alerts: Deserialized DeadlineAlert objects attached to the
task
+ :param task_instance: The TaskInstance that owns the deadlines
+ :param bundle_name: The bundle name for callback resolution
+ :param session: Database session
+ :return: The number of Deadline rows created
+ """
+ from airflow.serialization.definitions.deadline import
SerializedVariableInterval
+
+ created = 0
+ for deadline_alert in deadline_alerts:
+ interval = deadline_alert.interval
+ # The variable-backed interval is resolved lazily at evaluation time.
Every alert
+ # reaching the scheduler arrives decoded as a SerializedDeadlineAlert,
whose interval
+ # is a timedelta or SerializedVariableInterval.
+ if isinstance(interval, SerializedVariableInterval):
+ interval = interval.resolve()
+
+ deadline_time = deadline_alert.reference.evaluate_with(
+ session=session,
+ interval=interval,
+ dag_id=task_instance.dag_id,
+ run_id=task_instance.run_id,
+ )
+
+ if deadline_time is None:
+ continue
+
+ session.add(
+ Deadline(
+ deadline_time=deadline_time,
+ callback=deadline_alert.callback,
+ dagrun_id=task_instance.dag_run.id,
+ deadline_alert_id=None,
Review Comment:
`dagrun_id=task_instance.dag_run.id` will lazy-load the DagRun for each task
instance unless the relationship is already populated. In
`_process_task_deadline_alerts`, task instances come from
`DagRun.get_task_instances()` (query-based), so this becomes an N+1 query when
many tasks declare deadlines. Consider passing the DagRun id into
`create_deadlines_for_task_instance()` (or assigning `ti.dag_run = orm_dagrun`
at the call site) to avoid per-TI loads.
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -10106,6 +10107,233 @@ def test_process_expired_deadlines(self,
mock_handle_miss, session, dag_maker):
# Assert that all deadlines which are both expired and unhandled get
processed.
assert mock_handle_miss.call_count == 2
+ @mock.patch("airflow.models.Deadline.handle_miss")
+ def
test_process_expired_deadlines_skips_task_deadline_of_finished_in_time_task(
+ self, mock_handle_miss, session, dag_maker
+ ):
+ """A task-level deadline whose task instance finished before the
need-by time is pruned, not fired."""
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=1,
executors=[MockExecutor()])
+
+ now = timezone.utcnow()
+ callback_path = "classpath.notify"
+
+ dag_id = "test_deadline_dag_finished_in_time"
+ with dag_maker(dag_id=dag_id):
+ EmptyOperator(task_id="empty")
+ dagrun = dag_maker.create_dagrun()
+
+ task_instance = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == dagrun.run_id,
+ )
+ )
+ assert task_instance is not None
+ # The task succeeded well before the (now expired) need-by time.
+ task_instance.state = TaskInstanceState.SUCCESS
+ task_instance.end_date = now - timedelta(minutes=10)
+ session.flush()
+
+ deadline = Deadline(
+ deadline_time=now - timedelta(minutes=5),
+ callback=AsyncCallback(callback_path),
+ dagrun_id=dagrun.id,
+ dag_id=dag_id,
+ deadline_alert_id=None,
+ task_instance_id=task_instance.id,
+ )
+ session.add(deadline)
+ session.flush()
+ deadline_id = deadline.id
+
+ self.job_runner._execute()
+
+ # The deadline row was pruned instead of firing a false-alarm callback.
+ mock_handle_miss.assert_not_called()
+ assert session.get(Deadline, deadline_id) is None
+
+ @mock.patch("airflow.models.Deadline.handle_miss")
+ def
test_process_expired_deadlines_fires_task_deadline_of_still_running_task(
+ self, mock_handle_miss, session, dag_maker
+ ):
+ """A task-level deadline whose task is still running past the need-by
time fires normally."""
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=1,
executors=[MockExecutor()])
+
+ now = timezone.utcnow()
+ callback_path = "classpath.notify"
+
+ dag_id = "test_deadline_dag_still_running"
+ with dag_maker(dag_id=dag_id):
+ EmptyOperator(task_id="empty")
+ dagrun = dag_maker.create_dagrun()
+
+ task_instance = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == dagrun.run_id,
+ )
+ )
+ assert task_instance is not None
+ # The task started but has not finished as the need-by time passes.
+ task_instance.state = TaskInstanceState.RUNNING
+ task_instance.start_date = now - timedelta(minutes=10)
+ session.flush()
+
+ session.add(
+ Deadline(
+ deadline_time=now - timedelta(minutes=5),
+ callback=AsyncCallback(callback_path),
+ dagrun_id=dagrun.id,
+ dag_id=dag_id,
+ deadline_alert_id=None,
+ task_instance_id=task_instance.id,
+ )
+ )
+ session.flush()
+
+ self.job_runner._execute()
+
+ assert mock_handle_miss.call_count == 1
+
+ @mock.patch("airflow.models.Deadline.handle_miss")
+ def test_operator_declared_task_deadline_reaches_miss_loop(self,
mock_handle_miss, session, dag_maker):
+ """An operator-declared deadline materializes at DagRun creation and
fires in the miss loop."""
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=1,
executors=[MockExecutor()])
+
+ dag_id = "test_deadline_dag_operator_declared"
+ with dag_maker(dag_id=dag_id, serialized=True,
schedule=timedelta(days=1)):
+ EmptyOperator(
+ task_id="empty",
+ deadline=SdkDeadlineAlert(
+ reference=DeadlineReference.DAGRUN_LOGICAL_DATE,
+ interval=timedelta(minutes=-10),
+ callback=AsyncCallback("classpath.notify"),
+ ),
+ )
+ logical_date = timezone.utcnow() - timedelta(minutes=1)
+
+ # The deadline is anchored to the logical date, so pass a need-by time
already elapsed.
+ with time_machine.travel(logical_date, tick=False):
+ dagrun = dag_maker.create_dagrun(logical_date=logical_date)
+ session.commit()
+
+ task_instance = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == dagrun.run_id,
+ )
+ )
+ assert task_instance is not None
+ deadline = session.scalar(
+ select(Deadline).where(
+ Deadline.dagrun_id == dagrun.id,
+ Deadline.task_instance_id == task_instance.id,
+ )
+ )
+ # create_dagrun materialized a Deadline row owned by the task
instance, with the
+ # need-by time evaluated from the operator's DeadlineAlert
(logical date - 10min).
+ assert deadline is not None
+ assert deadline.deadline_time < timezone.utcnow()
+
+ self.job_runner._execute()
+
+ # The scheduler also creates a new scheduled DagRun for the dag; only
the expired
+ # deadline of the run under test fires.
+ # The scheduler also creates a new scheduled DagRun for the dag; the
deadline of the
+ # run under test is among the ones that fired. With the
instance-method mock, the
+ # Deadline arrives via the call's enclosing loop; match on the loop's
expired query.
Review Comment:
This block contains two overlapping comments explaining the same scheduler
behavior, which is noisy and a bit contradictory ("only" vs "among").
Consolidating to a single comment should improve readability.
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -10106,6 +10107,233 @@ def test_process_expired_deadlines(self,
mock_handle_miss, session, dag_maker):
# Assert that all deadlines which are both expired and unhandled get
processed.
assert mock_handle_miss.call_count == 2
+ @mock.patch("airflow.models.Deadline.handle_miss")
+ def
test_process_expired_deadlines_skips_task_deadline_of_finished_in_time_task(
+ self, mock_handle_miss, session, dag_maker
+ ):
+ """A task-level deadline whose task instance finished before the
need-by time is pruned, not fired."""
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=1,
executors=[MockExecutor()])
+
+ now = timezone.utcnow()
+ callback_path = "classpath.notify"
+
+ dag_id = "test_deadline_dag_finished_in_time"
+ with dag_maker(dag_id=dag_id):
+ EmptyOperator(task_id="empty")
+ dagrun = dag_maker.create_dagrun()
+
+ task_instance = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == dagrun.run_id,
+ )
+ )
+ assert task_instance is not None
+ # The task succeeded well before the (now expired) need-by time.
+ task_instance.state = TaskInstanceState.SUCCESS
+ task_instance.end_date = now - timedelta(minutes=10)
+ session.flush()
+
+ deadline = Deadline(
+ deadline_time=now - timedelta(minutes=5),
+ callback=AsyncCallback(callback_path),
+ dagrun_id=dagrun.id,
+ dag_id=dag_id,
+ deadline_alert_id=None,
+ task_instance_id=task_instance.id,
+ )
+ session.add(deadline)
+ session.flush()
+ deadline_id = deadline.id
+
+ self.job_runner._execute()
+
+ # The deadline row was pruned instead of firing a false-alarm callback.
+ mock_handle_miss.assert_not_called()
+ assert session.get(Deadline, deadline_id) is None
+
+ @mock.patch("airflow.models.Deadline.handle_miss")
+ def
test_process_expired_deadlines_fires_task_deadline_of_still_running_task(
+ self, mock_handle_miss, session, dag_maker
+ ):
+ """A task-level deadline whose task is still running past the need-by
time fires normally."""
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=1,
executors=[MockExecutor()])
+
+ now = timezone.utcnow()
+ callback_path = "classpath.notify"
+
+ dag_id = "test_deadline_dag_still_running"
+ with dag_maker(dag_id=dag_id):
+ EmptyOperator(task_id="empty")
+ dagrun = dag_maker.create_dagrun()
+
+ task_instance = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == dagrun.run_id,
+ )
+ )
+ assert task_instance is not None
+ # The task started but has not finished as the need-by time passes.
+ task_instance.state = TaskInstanceState.RUNNING
+ task_instance.start_date = now - timedelta(minutes=10)
+ session.flush()
+
+ session.add(
+ Deadline(
+ deadline_time=now - timedelta(minutes=5),
+ callback=AsyncCallback(callback_path),
+ dagrun_id=dagrun.id,
+ dag_id=dag_id,
+ deadline_alert_id=None,
+ task_instance_id=task_instance.id,
+ )
+ )
+ session.flush()
+
+ self.job_runner._execute()
+
+ assert mock_handle_miss.call_count == 1
+
+ @mock.patch("airflow.models.Deadline.handle_miss")
+ def test_operator_declared_task_deadline_reaches_miss_loop(self,
mock_handle_miss, session, dag_maker):
+ """An operator-declared deadline materializes at DagRun creation and
fires in the miss loop."""
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=1,
executors=[MockExecutor()])
+
+ dag_id = "test_deadline_dag_operator_declared"
+ with dag_maker(dag_id=dag_id, serialized=True,
schedule=timedelta(days=1)):
+ EmptyOperator(
+ task_id="empty",
+ deadline=SdkDeadlineAlert(
+ reference=DeadlineReference.DAGRUN_LOGICAL_DATE,
+ interval=timedelta(minutes=-10),
+ callback=AsyncCallback("classpath.notify"),
+ ),
+ )
+ logical_date = timezone.utcnow() - timedelta(minutes=1)
+
+ # The deadline is anchored to the logical date, so pass a need-by time
already elapsed.
+ with time_machine.travel(logical_date, tick=False):
+ dagrun = dag_maker.create_dagrun(logical_date=logical_date)
+ session.commit()
+
+ task_instance = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == dagrun.run_id,
+ )
+ )
+ assert task_instance is not None
+ deadline = session.scalar(
+ select(Deadline).where(
+ Deadline.dagrun_id == dagrun.id,
+ Deadline.task_instance_id == task_instance.id,
+ )
+ )
+ # create_dagrun materialized a Deadline row owned by the task
instance, with the
+ # need-by time evaluated from the operator's DeadlineAlert
(logical date - 10min).
+ assert deadline is not None
+ assert deadline.deadline_time < timezone.utcnow()
+
+ self.job_runner._execute()
+
+ # The scheduler also creates a new scheduled DagRun for the dag; only
the expired
+ # deadline of the run under test fires.
+ # The scheduler also creates a new scheduled DagRun for the dag; the
deadline of the
+ # run under test is among the ones that fired. With the
instance-method mock, the
+ # Deadline arrives via the call's enclosing loop; match on the loop's
expired query.
+ assert mock_handle_miss.call_count == 2
+
+ @mock.patch("airflow.models.Deadline.handle_miss")
+ def test_process_expired_deadlines_skips_task_deadline_of_removed_task(
+ self, mock_handle_miss, session, dag_maker
+ ):
+ """A deadline of a REMOVED (shrunk mapped) task instance is pruned: it
will never run."""
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=1,
executors=[MockExecutor()])
+
+ now = timezone.utcnow()
+ callback_path = "classpath.notify"
+
+ dag_id = "test_deadline_dag_removed_task"
+ with dag_maker(dag_id=dag_id):
+ EmptyOperator(task_id="empty")
+ dagrun = dag_maker.create_dagrun()
+
+ task_instance = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == dagrun.run_id,
+ )
+ )
+ assert task_instance is not None
+ task_instance.state = TaskInstanceState.REMOVED
+ session.flush()
+
+ deadline = Deadline(
+ deadline_time=now - timedelta(minutes=5),
+ callback=AsyncCallback(callback_path),
+ dagrun_id=dagrun.id,
+ dag_id=dag_id,
+ deadline_alert_id=None,
+ task_instance_id=task_instance.id,
+ )
+ session.add(deadline)
+ session.flush()
+ deadline_id = deadline.id
+
+ self.job_runner._execute()
+
+ mock_handle_miss.assert_not_called()
+ assert session.get(Deadline, deadline_id) is None
+
+ @mock.patch("airflow.models.Deadline.handle_miss")
+ def
test_process_expired_deadlines_skips_task_deadline_with_missing_task_instance(
+ self, mock_handle_miss, session, dag_maker
+ ):
Review Comment:
The test name says the deadline is "skipped" when the task instance
relationship is missing, but the assertion expects `handle_miss` to run
(call_count == 1). Renaming the test to match the behavior will make failures
easier to interpret.
--
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]
