Add optional debug spans for the scheduler loop

airflow-core/src/airflow/config_templates/config.yml

@@ -1516,9 +1516,6 @@ traces:
       description: |
         If True, then traces from Airflow internal methods are exported. Defaults to False.
       version_added: 3.1.0
-      version_deprecated: 3.2.0
-      deprecation_reason: |
-        This parameter is no longer used.
       type: string
       example: ~
       default: "False"

airflow-core/src/airflow/jobs/job.py

@@ -30,6 +30,7 @@
 from sqlalchemy.orm.session import make_transient
 
 from airflow._shared.observability.metrics import stats
+from airflow._shared.observability.traces import start_debug_span
 from airflow._shared.timezones import timezone
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
@@ -187,6 +188,7 @@ def on_kill(self):
         """Will be called when an external kill command is received."""
 
     @provide_session
+    @start_debug_span("job.heartbeat")
     def heartbeat(self, heartbeat_callback: Callable[..., None], *, session: Session = NEW_SESSION) -> None:
         """
         Update the job's entry in the database with the latest_heartbeat timestamp.

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

@@ -28,6 +28,7 @@
 import structlog
 from cachetools import TTLCache, cached
 from dateutil.relativedelta import relativedelta
+from opentelemetry import trace
 from sqlalchemy import (
     Boolean,
     Float,
@@ -56,6 +57,7 @@
 from sqlalchemy.sql import expression
 
 from airflow import settings
+from airflow._shared.observability.traces import start_debug_span
 from airflow._shared.timezones import timezone
 from airflow.assets.evaluation import AssetEvaluator
 from airflow.configuration import conf as airflow_conf
@@ -680,6 +682,7 @@ def deactivate_deleted_dags(
         return any_deactivated
 
     @classmethod
+    @start_debug_span("dag.dagmodel.dags_needing_dagruns")
     def dags_needing_dagruns(cls, session: Session) -> tuple[Any, dict[str, datetime]]:
         """
         Return (and lock) a list of Dag objects that are due to create a new DagRun.
@@ -804,6 +807,8 @@ def dag_ready(dag_id: str, cond: SerializedAssetBase, statuses: dict[UKey, bool]
             .limit(cls.NUM_DAGS_PER_DAGRUN_QUERY)
         )
 
+        trace.get_current_span().set_attribute("airflow.dag.asset_triggered_dags", len(triggered_date_by_dag))
+
         return (
             session.scalars(with_row_locks(query, of=cls, session=session, skip_locked=True)),
             triggered_date_by_dag,

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

@@ -78,6 +78,7 @@
     TRACE_SAMPLED_KEY,
     new_dagrun_trace_carrier,
     override_ids,
+    start_debug_span,
 )
 from airflow._shared.timezones import timezone
 from airflow.callbacks.callback_requests import DagCallbackRequest, DagRunContext
@@ -1241,6 +1242,7 @@ def _emit_dagrun_span(self, state: DagRunState):
             span.end()
 
     @provide_session
+    @start_debug_span("dagrun.update_state")
     def update_state(
         self, *, session: Session = NEW_SESSION, execute_callbacks: bool = True
     ) -> tuple[list[TI], DagCallbackRequest | None]:
@@ -1437,8 +1439,15 @@ def recalculate(self) -> _UnfinishedStates:
         session.merge(self)
         # We do not flush here for performance reasons(It increases queries count by +20)
 
+        update_state_span = trace.get_current_span()
+        update_state_span.set_attribute("airflow.dag_run.dag_id", self.dag_id)
+        update_state_span.set_attribute("airflow.dag_run.run_id", self.run_id)
+        update_state_span.set_attribute("airflow.dag_run.state", str(self.state))
+        update_state_span.set_attribute("airflow.dag_run.num_schedulable_tis", len(schedulable_tis))
+
         return schedulable_tis, callback
 
+    @start_debug_span("dagrun.task_instance_scheduling_decisions")
     @provide_session
     def task_instance_scheduling_decisions(self, *, session: Session = NEW_SESSION) -> TISchedulingDecision:
         tis = self.get_task_instances(session=session, state=State.task_states)
@@ -1481,6 +1490,12 @@ def _filter_tis_and_exclude_removed(dag: SerializedDAG, tis: list[TI]) -> Iterab
             schedulable_tis = []
             changed_tis = False
 
+        decision_span = trace.get_current_span()
+        decision_span.set_attribute("airflow.dag_run.num_tis", len(tis))
+        decision_span.set_attribute("airflow.dag_run.num_schedulable_tis", len(schedulable_tis))
+        decision_span.set_attribute("airflow.dag_run.num_unfinished_tis", len(unfinished_tis))
+        decision_span.set_attribute("airflow.dag_run.num_finished_tis", len(finished_tis))
+
         return TISchedulingDecision(
             tis=tis,
             schedulable_tis=schedulable_tis,
@@ -1618,6 +1633,7 @@ def execute_dag_callbacks(
                     ),
                 )
 
+    @start_debug_span("dagrun._get_ready_tis")
     def _get_ready_tis(
         self,
         schedulable_tis: list[TI],
@@ -1677,48 +1693,59 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None:
         expansion_happened = False
         # Set of task ids for which was already done _revise_map_indexes_if_mapped
         revised_map_index_task_ids: set[str] = set()
-        for schedulable in itertools.chain(schedulable_tis, additional_tis):
-            if TYPE_CHECKING:
-                assert isinstance(schedulable.task, Operator)
-            old_state = schedulable.state
-            if not schedulable.are_dependencies_met(session=session, dep_context=dep_context):
-                old_states[schedulable.key] = old_state
-                continue
-            # If schedulable is not yet expanded, try doing it now. This is
-            # called in two places: First and ideally in the mini scheduler at
-            # the end of LocalTaskJob, and then as an "expansion of last resort"
-            # in the scheduler to ensure that the mapped task is correctly
-            # expanded before executed. Also see _revise_map_indexes_if_mapped
-            # docstring for additional information.
-            new_tis = None
-            if schedulable.map_index < 0:
-                new_tis = _expand_mapped_task_if_needed(schedulable)
-                if new_tis is not None:
-                    additional_tis.extend(new_tis)
-                    expansion_happened = True
-            if new_tis is None and schedulable.state in SCHEDULEABLE_STATES:
-                # It's enough to revise map index once per task id,
-                # checking the map index for each mapped task significantly slows down scheduling
-                if schedulable.task.task_id not in revised_map_index_task_ids:
-                    ready_tis.extend(
-                        self._revise_map_indexes_if_mapped(
-                            schedulable.task, dag_version_id=schedulable.dag_version_id, session=session
-                        )
-                    )
-                    revised_map_index_task_ids.add(schedulable.task.task_id)
-
-                # _revise_map_indexes_if_mapped might mark the current task as REMOVED
-                # after calculating mapped task length, so we need to re-check
-                # the task state to ensure it's still schedulable
-                if schedulable.state in SCHEDULEABLE_STATES:
-                    ready_tis.append(schedulable)
+        with start_debug_span("dagrun.resolve_dependencies") as dep_span:
+            for schedulable in itertools.chain(schedulable_tis, additional_tis):
+                with start_debug_span("dagrun.resolve_dependencies_for_ti") as schedulable_span:
+                    if TYPE_CHECKING:
+                        assert isinstance(schedulable.task, Operator)
+                    schedulable_span.set_attribute("airflow.ti.task_id", schedulable.task.task_id)
+                    schedulable_span.set_attribute("airflow.ti.map_index", schedulable.map_index)
+                    old_state = schedulable.state
+                    if not schedulable.are_dependencies_met(session=session, dep_context=dep_context):
+                        old_states[schedulable.key] = old_state
+                        continue
+                    # If schedulable is not yet expanded, try doing it now. This is
+                    # called in two places: First and ideally in the mini scheduler at
+                    # the end of LocalTaskJob, and then as an "expansion of last resort"
+                    # in the scheduler to ensure that the mapped task is correctly
+                    # expanded before executed. Also see _revise_map_indexes_if_mapped
+                    # docstring for additional information.
+                    new_tis = None
+                    if schedulable.map_index < 0:
+                        new_tis = _expand_mapped_task_if_needed(schedulable)
+                        if new_tis is not None:
+                            additional_tis.extend(new_tis)
+                            expansion_happened = True
+                    if new_tis is None and schedulable.state in SCHEDULEABLE_STATES:
+                        # It's enough to revise map index once per task id,
+                        # checking the map index for each mapped task significantly slows down scheduling
+                        if schedulable.task.task_id not in revised_map_index_task_ids:
+                            ready_tis.extend(
+                                self._revise_map_indexes_if_mapped(
+                                    schedulable.task,
+                                    dag_version_id=schedulable.dag_version_id,
+                                    session=session,
+                                )
+                            )
+                            revised_map_index_task_ids.add(schedulable.task.task_id)
+
+                        # _revise_map_indexes_if_mapped might mark the current task as REMOVED
+                        # after calculating mapped task length, so we need to re-check
+                        # the task state to ensure it's still schedulable
+                        if schedulable.state in SCHEDULEABLE_STATES:
+                            ready_tis.append(schedulable)
+            dep_span.set_attribute(
+                "airflow.dag_run.num_dep_checks", len(schedulable_tis) + len(additional_tis)
+            )
 
         # Check if any ti changed state
         tis_filter = TI.filter_for_tis(old_states)
         if tis_filter is not None:
             fresh_tis = session.scalars(select(TI).where(tis_filter)).all()
             changed_tis = any(ti.state != old_states[ti.key] for ti in fresh_tis)
 
+        trace.get_current_span().set_attribute("airflow.dag_run.num_ready_tis", len(ready_tis))
+
         return ready_tis, changed_tis, expansion_happened
 
     def _are_premature_tis(
@@ -2159,6 +2186,7 @@ def get_latest_runs(cls, *, session: Session = NEW_SESSION) -> list[DagRun]:
             ).all()
         )
 
+    @start_debug_span("dagrun.schedule_tis")
     @provide_session
     def schedule_tis(
         self,
@@ -2295,6 +2323,8 @@ def schedule_tis(
                 )
                 count += getattr(result, "rowcount", 0)
 
+        trace.get_current_span().set_attribute("airflow.dag_run.num_scheduled_tis", count)
+
         return count
 
     @provide_session

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

@@ -74,6 +74,7 @@
     TASK_SPAN_DETAIL_LEVEL_KEY,
     new_dagrun_trace_carrier,
     new_task_run_carrier,
+    start_debug_span,
 )
 from airflow._shared.timezones import timezone
 from airflow.assets.manager import asset_manager
@@ -2135,13 +2136,18 @@ def _get_num_task_instances_of_state(
         return session.scalar(stmt) or 0
 
     @staticmethod
+    @start_debug_span("taskinstance.filter_for_tis")
     def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> ColumnElement[bool] | None:
         """Return SQLAlchemy filter to query selected task instances."""
         # DictKeys type, (what we often pass here from the scheduler) is not directly indexable :(
         # Or it might be a generator, but we need to be able to iterate over it more than once
         tis = list(tis)
 
+        filter_tis_span = trace.get_current_span()
+        filter_tis_span.set_attribute("airflow.taskinstance.filter_for_tis.num_tis", len(tis))
+
         if not tis:
+            filter_tis_span.set_attribute("airflow.taskinstance.filter_for_tis.result", "empty")
             return None
 
         first = tis[0]
@@ -2159,23 +2165,31 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> ColumnEleme
             map_indices.add(t.map_index)
             task_ids.add(t.task_id)
 
+        filter_tis_span.set_attribute("airflow.taskinstance.filter_for_tis.num_dag_ids", len(dag_ids))
+        filter_tis_span.set_attribute("airflow.taskinstance.filter_for_tis.num_run_ids", len(run_ids))
+        filter_tis_span.set_attribute("airflow.taskinstance.filter_for_tis.num_task_ids", len(task_ids))
+        filter_tis_span.set_attribute("airflow.taskinstance.filter_for_tis.num_map_indices", len(map_indices))
+
         # Common path optimisations: when all TIs are for the same dag_id and run_id, or same dag_id
         # and task_id -- this can be over 150x faster for huge numbers of TIs (20k+)
         if dag_ids == {dag_id} and run_ids == {run_id} and map_indices == {map_index}:
+            filter_tis_span.set_attribute("airflow.taskinstance.filter_for_tis.result", "same_dag_run_map")
             return and_(
                 TaskInstance.dag_id == dag_id,
                 TaskInstance.run_id == run_id,
                 TaskInstance.map_index == map_index,
                 TaskInstance.task_id.in_(task_ids),
             )
         if dag_ids == {dag_id} and task_ids == {first_task_id} and map_indices == {map_index}:
+            filter_tis_span.set_attribute("airflow.taskinstance.filter_for_tis.result", "same_dag_task_map")
             return and_(
                 TaskInstance.dag_id == dag_id,
                 TaskInstance.run_id.in_(run_ids),
                 TaskInstance.map_index == map_index,
                 TaskInstance.task_id == first_task_id,
             )
         if dag_ids == {dag_id} and run_ids == {run_id} and task_ids == {first_task_id}:
+            filter_tis_span.set_attribute("airflow.taskinstance.filter_for_tis.result", "same_dag_run_task")
             return and_(
                 TaskInstance.dag_id == dag_id,
                 TaskInstance.run_id == run_id,
@@ -2220,6 +2234,7 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> ColumnEleme
                         )
                     )
 
+        filter_tis_span.set_attribute("airflow.taskinstance.filter_for_tis.result", "grouped")
         return or_(*filter_condition)
 
     @classmethod

airflow-core/tests/integration/otel/dags/demo_dag.py

@@ -0,0 +1,73 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+
+from airflow import DAG
+from airflow.sdk import task
+
+logger = logging.getLogger("airflow.demo_dag")
+
+args = {
+    "owner": "airflow",
+    "start_date": datetime(2024, 9, 1),
+    "retries": 0,
+}
+
+
+@task
+def task1():
+    logger.info("task1 running")
+
+
+@task
+def task2():
+    logger.info("task2 running")
+
+
+@task
+def task3():
+    logger.info("task3 running")
+
+
+@task
+def task4():
+    logger.info("task4 running")
+
+
+@task
+def task5():
+    logger.info("task5 running")
+
+
+with DAG(
+    "demo_dag",
+    default_args=args,
+    schedule=None,
+    catchup=False,
+) as dag:
+    t1 = task1()
+    t2 = task2()
+    t3 = task3()
+    t4 = task4()
+    t5 = task5()
+
+    # task1 and task2 fan out in parallel, task3 joins on both,
+    # then task4 and task5 fan out in parallel.
+    [t1, t2] >> t3 >> [t4, t5]