From 4afffd3a1252d005e4dde6b2243103d03fbb2415 Mon Sep 17 00:00:00 2001 From: Christos Bisias Date: Wed, 1 Jul 2026 14:43:10 +0300 Subject: [PATCH] Add debug spans for the scheduler loop --- .../src/airflow/config_templates/config.yml | 3 - airflow-core/src/airflow/jobs/job.py | 2 + .../src/airflow/jobs/scheduler_job_runner.py | 315 ++++++++++++------ airflow-core/src/airflow/models/dag.py | 5 + airflow-core/src/airflow/models/dagrun.py | 100 ++++-- .../src/airflow/models/taskinstance.py | 15 + .../tests/integration/otel/dags/demo_dag.py | 73 ++++ .../tests/integration/otel/test_otel.py | 94 +++++- .../test_utils/integration_setup.py | 6 +- .../{otel_utils.py => otel_console_utils.py} | 37 +- .../test_utils/otel_jaeger_utils.py | 68 ++++ .../observability/traces/__init__.py | 25 +- ...el_utils.py => test_otel_console_utils.py} | 4 +- .../observability/test_otel_jaeger_utils.py | 173 ++++++++++ .../tests/observability/test_traces.py | 89 ++++- 15 files changed, 817 insertions(+), 192 deletions(-) create mode 100644 airflow-core/tests/integration/otel/dags/demo_dag.py rename devel-common/src/tests_common/test_utils/{otel_utils.py => otel_console_utils.py} (87%) create mode 100644 devel-common/src/tests_common/test_utils/otel_jaeger_utils.py rename shared/observability/tests/observability/{test_otel_utils.py => test_otel_console_utils.py} (99%) create mode 100644 shared/observability/tests/observability/test_otel_jaeger_utils.py diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index b3dc3f133d229..7a3389db7de5c 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/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" diff --git a/airflow-core/src/airflow/jobs/job.py b/airflow-core/src/airflow/jobs/job.py index b19cd103de908..f0cc92f935eb2 100644 --- a/airflow-core/src/airflow/jobs/job.py +++ b/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. diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index cb4c8f645520c..04c3529838046 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -34,6 +34,8 @@ from typing import TYPE_CHECKING, Any, cast from uuid import UUID +from opentelemetry import trace +from opentelemetry.context import context from sqlalchemy import ( CTE, Text, @@ -56,6 +58,7 @@ from airflow import settings from airflow._shared.observability.metrics import stats +from airflow._shared.observability.traces import start_debug_span from airflow._shared.timezones import timezone from airflow.api_fastapi.execution_api.datamodels.taskinstance import DagRun as DRDataModel, TIRunContext from airflow.assets.evaluation import AssetEvaluator @@ -410,6 +413,7 @@ def register_signals(self) -> ExitStack: return resetter + @start_debug_span("scheduler._get_team_names_for_dag_ids") def _get_team_names_for_dag_ids( self, dag_ids: Collection[str], session: Session ) -> dict[str, str | None]: @@ -424,6 +428,7 @@ def _get_team_names_for_dag_ids( :param session: Database session for queries :return: Dictionary mapping dag_id -> team_name (None if no team found) """ + trace.get_current_span().set_attribute("airflow.scheduler.num_dag_ids", len(dag_ids)) if not dag_ids: return {} @@ -543,6 +548,7 @@ def _debug_dump(self, signum: int, frame: FrameType | None) -> None: self.log.info("\n\t".join(map(repr, callstack))) self.log.info("-" * 80) + @start_debug_span("scheduler._executable_task_instances_to_queued") def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -> list[TI]: """ Find TIs that are ready for execution based on conditions. @@ -567,11 +573,13 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - # Optimization: to avoid littering the DB errors of "ERROR: canceling statement due to lock # timeout", try to take out a transactional advisory lock (unlocks automatically on # COMMIT/ROLLBACK) - lock_acquired = session.execute( - text("SELECT pg_try_advisory_xact_lock(:id)").bindparams( - id=DBLocks.SCHEDULER_CRITICAL_SECTION.value - ) - ).scalar() + with start_debug_span("scheduler.critical_section.advisory_lock") as advisory_lock_span: + lock_acquired = session.execute( + text("SELECT pg_try_advisory_xact_lock(:id)").bindparams( + id=DBLocks.SCHEDULER_CRITICAL_SECTION.value + ) + ).scalar() + advisory_lock_span.set_attribute("airflow.scheduler.lock_acquired", bool(lock_acquired)) if lock_acquired is None: lock_acquired = False if not lock_acquired: @@ -724,8 +732,12 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - timer.start() try: - locked_query = with_row_locks(query, of=TI, session=session, skip_locked=True) - task_instances_to_examine = session.scalars(locked_query).all() + with start_debug_span("scheduler.critical_section.task_instances_to_examine") as select_span: + locked_query = with_row_locks(query, of=TI, session=session, skip_locked=True) + task_instances_to_examine = session.scalars(locked_query).all() + select_span.set_attribute( + "airflow.scheduler.tis_examined", len(task_instances_to_examine) + ) if self.log.isEnabledFor(logging.DEBUG): self.log.debug("Length of the tis to examine is %d", len(task_instances_to_examine)) @@ -1057,33 +1069,37 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - .execution_options(synchronize_session=False) ) - if pre_assign_executors: - # Read the DB-generated UUIDs back onto the in-memory objects so the - # workload DTO carries them through to send_workload_to_executor (the - # objects are about to be detached by make_transient). Use RETURNING - # where supported (PostgreSQL); fall back to a SELECT for MySQL and - # SQLite (RETURNING requires SQLite 3.35+ which isn't guaranteed). - if get_dialect_name(session) == "postgresql": - result = session.execute(queued_update.returning(TI.id, TI.external_executor_id)) - id_map = {row[0]: row[1] for row in result} + with start_debug_span("scheduler.critical_section.update_tis_to_queued") as update_span: + update_span.set_attribute("airflow.scheduler.tis_to_update", len(executable_tis)) + if pre_assign_executors: + # Read the DB-generated UUIDs back onto the in-memory objects so the + # workload DTO carries them through to send_workload_to_executor (the + # objects are about to be detached by make_transient). Use RETURNING + # where supported (PostgreSQL); fall back to a SELECT for MySQL and + # SQLite (RETURNING requires SQLite 3.35+ which isn't guaranteed). + if get_dialect_name(session) == "postgresql": + result = session.execute(queued_update.returning(TI.id, TI.external_executor_id)) + id_map = {row[0]: row[1] for row in result} + else: + session.execute(queued_update) + id_rows = session.execute( + select(TI.id, TI.external_executor_id).where(filter_for_tis) + ).all() + id_map = {row[0]: row[1] for row in id_rows} + for ti in executable_tis: + ti.external_executor_id = id_map.get(ti.id) else: session.execute(queued_update) - id_rows = session.execute( - select(TI.id, TI.external_executor_id).where(filter_for_tis) - ).all() - id_map = {row[0]: row[1] for row in id_rows} - for ti in executable_tis: - ti.external_executor_id = id_map.get(ti.id) - else: - session.execute(queued_update) for ti in executable_tis: ti.emit_state_change_metric(TaskInstanceState.QUEUED) for ti in executable_tis: make_transient(ti) + trace.get_current_span().set_attribute("airflow.scheduler.num_executable_tis", len(executable_tis)) return executable_tis + @start_debug_span("scheduler._enqueue_task_instances_with_queued_state") def _enqueue_task_instances_with_queued_state( self, task_instances: list[TI], executor: BaseExecutor, session: Session ) -> None: @@ -1094,6 +1110,9 @@ def _enqueue_task_instances_with_queued_state( :param executor: The executor to enqueue tasks for :param session: The session object """ + enqueue_span = trace.get_current_span() + enqueue_span.set_attribute("airflow.executor.class", type(executor).__name__) + enqueue_span.set_attribute("airflow.scheduler.num_tis_to_enqueue", len(task_instances)) def _get_sentry_integration(executor: BaseExecutor) -> str: try: @@ -1114,32 +1133,39 @@ def _get_sentry_integration(executor: BaseExecutor) -> str: # actually enqueue them for ti in task_instances: - if ti.dag_run.state in State.finished_dr_states: - ti.set_state(None, session=session) - continue - if not ti.dag_version_id: - self.log.warning( - "TaskInstance %s does not have a dag_version_id set, cannot be enqueued. " - "This would get unstuck and dag_version_id updated.", + with start_debug_span("scheduler.enqueue_ti") as enqueue_ti_span: + enqueue_ti_span.set_attribute("airflow.ti.dag_id", ti.dag_id) + enqueue_ti_span.set_attribute("airflow.ti.run_id", ti.run_id) + enqueue_ti_span.set_attribute("airflow.ti.task_id", ti.task_id) + enqueue_ti_span.set_attribute("airflow.ti.map_index", ti.map_index) + enqueue_ti_span.set_attribute("airflow.ti.try_number", ti.try_number) + if ti.dag_run.state in State.finished_dr_states: + ti.set_state(None, session=session) + continue + if not ti.dag_version_id: + self.log.warning( + "TaskInstance %s does not have a dag_version_id set, cannot be enqueued. " + "This would get unstuck and dag_version_id updated.", + ti, + ) + continue + + self.log.debug( + "Queueing workload for TI: %s try_number=%d state=%s scheduler_job_id=%s executor=%s", ti, + ti.try_number, + ti.state, + self.job.id, + executor, ) - continue - - self.log.debug( - "Queueing workload for TI: %s try_number=%d state=%s scheduler_job_id=%s executor=%s", - ti, - ti.try_number, - ti.state, - self.job.id, - executor, - ) - workload = workloads.ExecuteTask.make( - ti, - generator=executor.jwt_generator, - sentry_integration=_get_sentry_integration(executor), - ) - executor.queue_workload(workload, session=session) + workload = workloads.ExecuteTask.make( + ti, + generator=executor.jwt_generator, + sentry_integration=_get_sentry_integration(executor), + ) + executor.queue_workload(workload, session=session) + @start_debug_span("scheduler._critical_section_enqueue_task_instances") def _critical_section_enqueue_task_instances(self, session: Session) -> int: """ Enqueues TaskInstances for execution. @@ -1189,8 +1215,10 @@ def _critical_section_enqueue_task_instances(self, session: Session) -> int: self._enqueue_task_instances_with_queued_state(queued_tis_per_executor, executor, session=session) + trace.get_current_span().set_attribute("airflow.scheduler.num_queued_tis", len(queued_tis)) return len(queued_tis) + @start_debug_span("scheduler._enqueue_executor_callbacks") def _enqueue_executor_callbacks(self, session: Session) -> None: """ Enqueue ExecutorCallback workloads to executors. @@ -1214,6 +1242,7 @@ def _enqueue_executor_callbacks(self, session: Session) -> None: .order_by(ExecutorCallback.priority_weight.desc()) .limit(max_callbacks) ).all() + trace.get_current_span().set_attribute("airflow.scheduler.pending_callbacks", len(pending_callbacks)) if not pending_callbacks: return @@ -1261,7 +1290,9 @@ def _enqueue_executor_callbacks(self, session: Session) -> None: session.add(callback) @staticmethod + @start_debug_span("scheduler._process_task_event_logs") def _process_task_event_logs(log_records: deque[Log], session: Session): + trace.get_current_span().set_attribute("airflow.scheduler.num_task_event_logs", len(log_records)) objects = (log_records.popleft() for _ in range(len(log_records))) session.bulk_save_objects(objects=objects, preserve_order=False) @@ -1279,15 +1310,18 @@ def _is_metrics_enabled(): def _is_tracing_enabled(): return conf.getboolean("traces", "otel_on") + @start_debug_span("scheduler._process_executor_events") def _process_executor_events(self, executor: BaseExecutor, session: Session) -> int: try: - return SchedulerJobRunner.process_executor_events( + events_processed = SchedulerJobRunner.process_executor_events( executor=executor, job_id=self.job.id, scheduler_dag_bag=self.scheduler_dag_bag, session=session, eagerly_load_dag_tags=self._dag_tags_in_metrics, ) + trace.get_current_span().set_attribute("airflow.scheduler.events_processed", events_processed) + return events_processed except Exception as exc: stats.incr("scheduler.executor_events.failed", tags={"exception_class": type(exc).__name__}) raise @@ -1831,7 +1865,14 @@ def _run_scheduler_loop(self) -> None: # Reset per-loop team name cache so changes to bundle-team assignments # are picked up each iteration without requiring a scheduler restart. self._dag_id_to_team_name = {} - with stats.timer("scheduler.scheduler_loop_duration") as timer: + with ( + start_debug_span( + "scheduler.scheduler_loop", + context=context.Context(), + attributes={"airflow.scheduler.loop_count": loop_count}, + ) as loop_iteration_span, + stats.timer("scheduler.scheduler_loop_duration") as timer, + ): with create_session() as session: # This will schedule for as many executors as possible. num_queued_tis = self._do_scheduling(session) @@ -1842,9 +1883,15 @@ def _run_scheduler_loop(self) -> None: # either a no-op, or they will check-in on currently running tasks and send out new # events to be processed below. for executor in self.executors: - with stats.timer( - "scheduler.executor_heartbeat_duration", - tags={"executor": type(executor).__name__}, + with ( + start_debug_span( + "scheduler.executor.heartbeat", + attributes={"airflow.executor.class": type(executor).__name__}, + ), + stats.timer( + "scheduler.executor_heartbeat_duration", + tags={"executor": type(executor).__name__}, + ), ): executor.heartbeat() @@ -1872,16 +1919,18 @@ def _run_scheduler_loop(self) -> None: .where(~Deadline.missed) .options(selectinload(Deadline.callback), selectinload(Deadline.dagrun)) ) - for deadline in session.scalars( - with_row_locks( - deadline_query, - of=Deadline, - session=session, - skip_locked=True, - key_share=False, - ) - ): - deadline.handle_miss(session) + with start_debug_span("scheduler.handle_missed_deadlines"): + for deadline in session.scalars( + with_row_locks( + deadline_query, + of=Deadline, + session=session, + skip_locked=True, + key_share=False, + ) + ): + with start_debug_span("scheduler.handle_deadline_miss"): + deadline.handle_miss(session) # Route ExecutorCallback workloads to executors (similar to task routing) self._enqueue_executor_callbacks(session) @@ -1894,12 +1943,19 @@ def _run_scheduler_loop(self) -> None: ) # Run any pending timed events - next_event = timers.run(blocking=False) - self.log.debug("Next timed event is in %f", next_event) + with start_debug_span("scheduler.timers.run"): + next_event = timers.run(blocking=False) + self.log.debug("Next timed event is in %f", next_event) + + loop_iteration_span.set_attribute("airflow.scheduler.num_queued_tis", num_queued_tis) + loop_iteration_span.set_attribute( + "airflow.scheduler.num_finished_events", num_finished_events + ) + idle_in_this_run = not num_queued_tis and not num_finished_events + loop_iteration_span.set_attribute("airflow.scheduler.loop_iteration.idle", idle_in_this_run) self.log.debug("Ran scheduling loop in %.2f ms", timer.duration) - idle_in_this_run = not num_queued_tis and not num_finished_events if not is_unit_test and idle_in_this_run: # If the scheduler is doing things, don't sleep. This means when there is work to do, the # scheduler will run "as quick as possible", but when it's stopped, it can sleep, dropping CPU @@ -1921,6 +1977,7 @@ def _run_scheduler_loop(self) -> None: ) break + @start_debug_span("scheduler._do_scheduling") def _do_scheduling(self, session: Session) -> int: """ Make the main scheduling decisions. @@ -1963,11 +2020,15 @@ def _do_scheduling(self, session: Session) -> int: # examining, rather than making one query per DagRun. # Materialize into a list because the multi-team block below iterates # the result and ScalarResult is a one-pass iterator. - dag_runs = list( - DagRun.get_running_dag_runs_to_examine( - session=session, eagerly_load_dag_tags=self._dag_tags_in_metrics + with start_debug_span("dagrun.get_running_dag_runs_to_examine") as dag_runs_span: + dag_runs = list( + DagRun.get_running_dag_runs_to_examine( + session=session, eagerly_load_dag_tags=self._dag_tags_in_metrics + ) ) - ) + # This span isn't added as a method decorator so that + # the below attribute can also be added as a tag. + dag_runs_span.set_attribute("airflow.scheduler.dagrun.dag_run_count", len(dag_runs)) if self._multi_team and dag_runs: unique_dag_ids = {dr.dag_id for dr in dag_runs} @@ -1991,13 +2052,19 @@ def _do_scheduling(self, session: Session) -> int: else: self.log.error("DAG '%s' not found in serialized_dag table", dag_run.dag_id) - with prohibit_commit(session) as guard: + with ( + prohibit_commit(session) as guard, + start_debug_span("scheduler.critical_section") as critical_section_span, + ): # Without this, the session has an invalid view of the DB session.expunge_all() # END: schedule TIs # Attempt to schedule even if some executors are full but not all. total_free_executor_slots = sum([executor.slots_available for executor in self.executors]) + critical_section_span.set_attribute( + "airflow.scheduler.free_executor_slots", total_free_executor_slots + ) if total_free_executor_slots <= 0: # We know we can't do anything here, so don't even try! self.log.debug("All executors are full, skipping critical section") @@ -2013,12 +2080,15 @@ def _do_scheduling(self, session: Session) -> int: # Make sure we only sent this metric if we obtained the lock, otherwise we'll skew the # metric, way down timer.stop(send=True) + critical_section_span.set_attribute("airflow.scheduler.lock_acquired", True) + critical_section_span.set_attribute("airflow.scheduler.num_queued_tis", num_queued_tis) except OperationalError as e: timer.stop(send=False) if is_lock_not_available_error(error=e): self.log.debug("Critical section lock held by another Scheduler") stats.incr("scheduler.critical_section_busy") + critical_section_span.set_attribute("airflow.scheduler.lock_acquired", False) session.rollback() return 0 raise @@ -2400,12 +2470,16 @@ def _create_dagruns_for_partitioned_asset_dags(self, session: Session) -> set[st return partition_dag_ids @retry_db_transaction + @start_debug_span("scheduler._create_dagruns_for_dags") def _create_dagruns_for_dags(self, guard: CommitProhibitorGuard, session: Session) -> None: """Find Dag Models needing DagRuns and Create Dag Runs with retries in case of OperationalError.""" partition_dag_ids: set[str] = self._create_dagruns_for_partitioned_asset_dags(session) query, triggered_date_by_dag = DagModel.dags_needing_dagruns(session) all_dags_needing_dag_runs = set(query.all()) + trace.get_current_span().set_attribute( + "airflow.scheduler.dags_needing_runs", len(all_dags_needing_dag_runs) + ) asset_triggered_dags = [d for d in all_dags_needing_dag_runs if d.dag_id in triggered_date_by_dag] non_asset_dags = { d @@ -2456,8 +2530,10 @@ def _mark_backfills_complete(self, *, session: Session = NEW_SESSION) -> None: for b in backfills: b.completed_at = now + @start_debug_span("scheduler._create_dag_runs") def _create_dag_runs(self, dag_models: Collection[DagModel], session: Session) -> None: """Create a DAG run and update the dag_model to control if/when the next DAGRun should be created.""" + trace.get_current_span().set_attribute("airflow.scheduler.num_dag_models", len(dag_models)) # Bulk Fetch DagRuns with dag_id and logical_date same # as DagModel.dag_id and DagModel.next_dagrun # This list is used to verify if the DagRun already exist so that we don't attempt to create @@ -2592,6 +2668,7 @@ def _create_dag_runs(self, dag_models: Collection[DagModel], session: Session) - # TODO[HA]: Should we do a session.flush() so we don't have to keep lots of state/object in # memory for larger dags? or expunge_all() + @start_debug_span("scheduler._create_dag_runs_asset_triggered") def _create_dag_runs_asset_triggered( self, *, @@ -2599,6 +2676,7 @@ def _create_dag_runs_asset_triggered( session: Session, ) -> None: """For Dags that are triggered by assets, create Dag runs.""" + trace.get_current_span().set_attribute("airflow.scheduler.num_asset_triggered_dags", len(dag_models)) for dag_model in dag_models: dag = self._get_current_dag(dag_id=dag_model.dag_id, session=session) if not dag: @@ -2753,9 +2831,11 @@ def _lock_backfills(self, dag_runs: Collection[DagRun], session: Session) -> dic return locked_backfills + @start_debug_span("scheduler._start_queued_dagruns") def _start_queued_dagruns(self, session: Session) -> None: """Find DagRuns in queued state and decide moving them to running state.""" dag_runs: Collection[DagRun] = list(DagRun.get_queued_dag_runs_to_set_running(session)) + trace.get_current_span().set_attribute("airflow.scheduler.queued_dag_runs", len(dag_runs)) # Lock backfills to prevent race conditions with concurrent schedulers locked_backfills = self._lock_backfills(dag_runs, session) @@ -2806,46 +2886,54 @@ def _update_state(dag: SerializedDAG, dag_run: DagRun): dag_id = dag_run.dag_id run_id = dag_run.run_id backfill_id = dag_run.backfill_id - dag = dag_run.dag = cached_get_dag(dag_run) - if not dag: - self.log.error("DAG '%s' not found in serialized_dag table", dag_run.dag_id) - continue - active_runs = active_runs_of_dags[(dag_id, backfill_id)] - if backfill_id is not None: - if backfill_id not in locked_backfills: - # Another scheduler has this backfill locked, skip this run - continue - backfill = dag_run.backfill - if active_runs >= backfill.max_active_runs: - # todo: delete all "candidate dag runs" from list for this dag right now - self.log.info( - "dag cannot be started due to backfill max_active_runs constraint; " - "active_runs=%s max_active_runs=%s dag_id=%s run_id=%s", - active_runs, - backfill.max_active_runs, - dag_id, - run_id, - ) - continue - elif dag_run.max_active_runs: - # Using dag_run.max_active_runs which links to DagModel to ensure we are checking - # against the most recent changes on the dag and not using stale serialized dag - if active_runs >= dag_run.max_active_runs: - # todo: delete all candidate dag runs for this dag from list right now - self.log.info( - "dag cannot be started due to dag max_active_runs constraint; " - "active_runs=%s max_active_runs=%s dag_id=%s run_id=%s", - active_runs, - dag_run.max_active_runs, - dag_run.dag_id, - dag_run.run_id, - ) + with start_debug_span("scheduler.update_dag_run_state_to_running") as update_state_span: + update_state_span.set_attribute("airflow.dag_run.dag_id", dag_id) + update_state_span.set_attribute("airflow.dag_run.run_id", run_id) + if backfill_id is not None: + update_state_span.set_attribute("airflow.dag_run.backfill_id", backfill_id) + with start_debug_span("scheduler.cached_get_dag") as cached_get_dag_span: + cached_get_dag_span.set_attribute("airflow.dag.dag_id", dag_id) + dag = dag_run.dag = cached_get_dag(dag_run) + if not dag: + self.log.error("DAG '%s' not found in serialized_dag table", dag_run.dag_id) continue - active_runs_of_dags[(dag_run.dag_id, backfill_id)] += 1 - _update_state(dag, dag_run) - dag_run.notify_dagrun_state_changed(msg="started") + active_runs = active_runs_of_dags[(dag_id, backfill_id)] + if backfill_id is not None: + if backfill_id not in locked_backfills: + # Another scheduler has this backfill locked, skip this run + continue + backfill = dag_run.backfill + if active_runs >= backfill.max_active_runs: + # todo: delete all "candidate dag runs" from list for this dag right now + self.log.info( + "dag cannot be started due to backfill max_active_runs constraint; " + "active_runs=%s max_active_runs=%s dag_id=%s run_id=%s", + active_runs, + backfill.max_active_runs, + dag_id, + run_id, + ) + continue + elif dag_run.max_active_runs: + # Using dag_run.max_active_runs which links to DagModel to ensure we are checking + # against the most recent changes on the dag and not using stale serialized dag + if active_runs >= dag_run.max_active_runs: + # todo: delete all candidate dag runs for this dag from list right now + self.log.info( + "dag cannot be started due to dag max_active_runs constraint; " + "active_runs=%s max_active_runs=%s dag_id=%s run_id=%s", + active_runs, + dag_run.max_active_runs, + dag_run.dag_id, + dag_run.run_id, + ) + continue + active_runs_of_dags[(dag_run.dag_id, backfill_id)] += 1 + _update_state(dag, dag_run) + dag_run.notify_dagrun_state_changed(msg="started") @retry_db_transaction + @start_debug_span("scheduler._schedule_all_dag_runs") def _schedule_all_dag_runs( self, guard: CommitProhibitorGuard, @@ -2863,8 +2951,10 @@ def _schedule_all_dag_runs( except Exception: self.log.exception("Error scheduling DAG run %s of %s", run.run_id, run.dag_id) guard.commit() + trace.get_current_span().set_attribute("airflow.scheduler.dag_runs_scheduled", len(callback_tuples)) return callback_tuples + @start_debug_span("scheduler._schedule_dag_run") def _schedule_dag_run( self, dag_run: DagRun, @@ -2876,6 +2966,10 @@ def _schedule_dag_run( :param dag_run: The DagRun to schedule :return: Callback that needs to be executed """ + current_span = trace.get_current_span() + current_span.set_attribute("airflow.dag_run.dag_id", dag_run.dag_id) + current_span.set_attribute("airflow.dag_run.run_id", dag_run.run_id) + callback: DagCallbackRequest | None = None dag = dag_run.dag = self.scheduler_dag_bag.get_dag_for_run(dag_run=dag_run, session=session) @@ -2972,6 +3066,8 @@ def _schedule_dag_run( # TODO[HA]: Rename update_state -> schedule_dag_run, ?? something else? schedulable_tis, callback_to_run = dag_run.update_state(session=session, execute_callbacks=False) + current_span.set_attribute("airflow.scheduler.num_schedulable_tis", len(schedulable_tis)) + current_span.set_attribute("airflow.dag_run.state", str(dag_run.state)) if dag_run.state in State.finished_dr_states and dag_run.run_type in ( DagRunType.SCHEDULED, @@ -3929,6 +4025,7 @@ def _cleanup_orphaned_asset_state_store(*, session: Session) -> None: delete(AssetStateStoreModel).where(AssetStateStoreModel.asset_id.not_in(active_asset_ids)) ) + @start_debug_span("scheduler._enqueue_connection_tests") def _enqueue_connection_tests(self, *, session: Session) -> None: """ Enqueue pending connection tests to executors that support them. @@ -3960,6 +4057,9 @@ def _enqueue_connection_tests(self, *, session: Session) -> None: ) stats.gauge("connection_test.active", active_count) stats.gauge("connection_test.pending", pending_count) + connection_test_span = trace.get_current_span() + connection_test_span.set_attribute("airflow.scheduler.connection_tests_active", active_count) + connection_test_span.set_attribute("airflow.scheduler.connection_tests_pending", pending_count) budget = max_concurrency - active_count if budget <= 0: @@ -4055,6 +4155,7 @@ def _reap_stale_connection_tests(self, *, session: Session = NEW_SESSION) -> Non session.flush() + @start_debug_span("scheduler._executor_to_workloads") def _executor_to_workloads( self, workloads: Iterable[SchedulerWorkload], diff --git a/airflow-core/src/airflow/models/dag.py b/airflow-core/src/airflow/models/dag.py index f827785610446..e166a16973d3c 100644 --- a/airflow-core/src/airflow/models/dag.py +++ b/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, diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index d6610b3663e7e..6c6d30aa31509 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/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,41 +1693,50 @@ 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) @@ -1719,6 +1744,8 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | 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 diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 0dd10507fc863..3a2deb1c649b7 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/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,9 +2165,15 @@ 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, @@ -2169,6 +2181,7 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> ColumnEleme 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), @@ -2176,6 +2189,7 @@ def filter_for_tis(tis: Iterable[TaskInstance | TaskInstanceKey]) -> ColumnEleme 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 diff --git a/airflow-core/tests/integration/otel/dags/demo_dag.py b/airflow-core/tests/integration/otel/dags/demo_dag.py new file mode 100644 index 0000000000000..ced29d850affe --- /dev/null +++ b/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] diff --git a/airflow-core/tests/integration/otel/test_otel.py b/airflow-core/tests/integration/otel/test_otel.py index 39d7ee6bdab61..613f578fd1f9d 100644 --- a/airflow-core/tests/integration/otel/test_otel.py +++ b/airflow-core/tests/integration/otel/test_otel.py @@ -38,9 +38,10 @@ unpause_trigger_dag_and_get_run_id, wait_for_dag_run, ) -from tests_common.test_utils.otel_utils import ( - dump_airflow_metadata_db, - extract_metrics_from_output, +from tests_common.test_utils.otel_console_utils import extract_metrics_from_output +from tests_common.test_utils.otel_jaeger_utils import ( + get_span_tags, + provided_child_spans_found_under_span, ) if TYPE_CHECKING: @@ -111,6 +112,20 @@ def print_ti_output_for_dag_run(dag_id: str, run_id: str): print("\n===== END =====\n") +def get_non_idle_loop_trace(all_traces: list[dict]) -> dict | None: + for trace in all_traces: + scheduler_loop_span = next( + (s for s in trace["spans"] if s["operationName"] == "scheduler.scheduler_loop"), + None, + ) + if scheduler_loop_span is None: + continue + # A non-idle scheduler_loop span means this iteration actually scheduled something. + if get_span_tags(scheduler_loop_span).get("airflow.scheduler.loop_iteration.idle") is False: + return trace + return None + + @pytest.mark.integration("otel") @pytest.mark.backend("postgres") class TestOtelIntegration: @@ -183,7 +198,7 @@ def setup_class(cls): migrate_command = ["airflow", "db", "migrate"] subprocess.run(migrate_command, check=True, env=os.environ.copy()) - cls.dags = serialize_and_get_dags(dag_folder=cls.dag_folder) + cls.dags = serialize_and_get_dags(dag_folder=cls.dag_folder, num_of_dags=2) def dag_execution_for_testing_metrics(self, capfd): # Metrics. @@ -388,10 +403,6 @@ def test_dag_execution_succeeds(self, capfd, task_span_detail_level, expected_hi print_ti_output_for_dag_run(dag_id=dag_id, run_id=run_id) finally: - if self.log_level == "debug": - with create_session() as session: - dump_airflow_metadata_db(session) - # Terminate the processes. terminate_process(scheduler_process) scheduler_status = scheduler_process.poll() @@ -440,3 +451,70 @@ def get_parent_span_id(span): nested = get_span_hierarchy() assert nested == expected_hierarchy + + @pytest.mark.execution_timeout(160) + def test_scheduler_debug_traces(self, capfd): + # Enable debug traces. + os.environ["AIRFLOW__TRACES__OTEL_DEBUG_TRACES_ON"] = "True" + + scheduler_process = None + apiserver_process = None + try: + # Start the processes here and not as fixtures or in a common setup, + # so that the test can capture their output. + scheduler_process, apiserver_process = start_scheduler() + + dag_id = "demo_dag" + + assert len(self.dags) > 0 + dag = self.dags[dag_id] + + assert dag is not None + + conf = None + + run_id_1 = unpause_trigger_dag_and_get_run_id(dag_id=dag_id, conf=conf) + run_id_2 = unpause_trigger_dag_and_get_run_id(dag_id=dag_id, conf=conf) + + wait_for_dag_run(dag_id=dag_id, run_id=run_id_1, max_wait_time=90) + wait_for_dag_run(dag_id=dag_id, run_id=run_id_2, max_wait_time=90) + + time.sleep(10) + + print_ti_output_for_dag_run(dag_id=dag_id, run_id=run_id_1) + print_ti_output_for_dag_run(dag_id=dag_id, run_id=run_id_2) + finally: + # Terminate the processes. + terminate_process(scheduler_process) + scheduler_status = scheduler_process.poll() + assert scheduler_status is not None, ( + "The scheduler_1 process status is None, which means that it hasn't terminated as expected." + ) + + terminate_process(apiserver_process) + apiserver_status = apiserver_process.poll() + assert apiserver_status is not None, ( + "The apiserver process status is None, which means that it hasn't terminated as expected." + ) + + capfd.readouterr() + + host = "jaeger" + service_name = os.environ.get("OTEL_SERVICE_NAME", "test") + r = requests.get(f"http://{host}:16686/api/traces?service={service_name}") + traces = r.json()["data"] + + # There are spans that don't get exported on every loop iteration. The ones below + # reliably re-appear every time that the loop does some work. + expected_children_spans = [ + "scheduler._do_scheduling", + "scheduler.critical_section", + "scheduler._process_executor_events", + ] + + non_idle_trace = get_non_idle_loop_trace(traces) + assert non_idle_trace is not None, "expected at least one non-idle scheduler_loop span" + + assert provided_child_spans_found_under_span( + non_idle_trace, "scheduler.scheduler_loop", expected_children_spans + ), f"expected {expected_children_spans} as descendants of scheduler.scheduler_loop" diff --git a/devel-common/src/tests_common/test_utils/integration_setup.py b/devel-common/src/tests_common/test_utils/integration_setup.py index 59476d8c4d06b..4984a18f43804 100644 --- a/devel-common/src/tests_common/test_utils/integration_setup.py +++ b/devel-common/src/tests_common/test_utils/integration_setup.py @@ -109,12 +109,16 @@ def start_scheduler(capture_output: bool = False): return scheduler_process, apiserver_process -def serialize_and_get_dags(dag_folder) -> dict[str, SerializedDAG]: +def serialize_and_get_dags(dag_folder: str, num_of_dags: int | None = None) -> dict[str, SerializedDAG]: log.info("Serializing Dags from directory %s", dag_folder) # Load DAGs from the dag directory. dag_bag = DagBag(dag_folder=dag_folder) dag_ids = dag_bag.dag_ids + if num_of_dags is None: + assert len(dag_ids) > 0 + else: + assert len(dag_ids) == num_of_dags dag_dict: dict[str, SerializedDAG] = {} with create_session() as session: diff --git a/devel-common/src/tests_common/test_utils/otel_utils.py b/devel-common/src/tests_common/test_utils/otel_console_utils.py similarity index 87% rename from devel-common/src/tests_common/test_utils/otel_utils.py rename to devel-common/src/tests_common/test_utils/otel_console_utils.py index e540616092262..0dd332561cca7 100644 --- a/devel-common/src/tests_common/test_utils/otel_utils.py +++ b/devel-common/src/tests_common/test_utils/otel_console_utils.py @@ -14,47 +14,16 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +"""Helpers for spans and metrics captured from the OTel console (stdout) exporter.""" + from __future__ import annotations import json import logging -import pprint from collections import defaultdict from typing import Literal -from sqlalchemy import inspect, select - -log = logging.getLogger("integration.otel.otel_utils") - - -def dump_airflow_metadata_db(session): - from airflow.models import Base - - inspector = inspect(session.bind) - all_tables = inspector.get_table_names() - - # dump with the entire db - db_dump = {} - - log.debug("\n-----START_airflow_db_dump-----\n") - - for table_name in all_tables: - log.debug("\nDumping table: %s", table_name) - table = Base.metadata.tables.get(table_name) - if table is not None: - results = [dict(row._mapping) for row in session.execute(select(table)).all()] - db_dump[table_name] = results - # Pretty-print the table contents - if table_name == "connection": - filtered_results = [row for row in results if row.get("conn_id") == "airflow_db"] - pprint.pprint({table_name: filtered_results}, width=120) - else: - pprint.pprint({table_name: results}, width=120) - else: - log.debug("Table %s not found in metadata.", table_name) - - log.debug("\nAirflow metadata database dump complete.") - log.debug("\n-----END_airflow_db_dump-----\n") +log = logging.getLogger("integration.otel.otel_console_utils") def extract_task_event_value(line: str) -> str: diff --git a/devel-common/src/tests_common/test_utils/otel_jaeger_utils.py b/devel-common/src/tests_common/test_utils/otel_jaeger_utils.py new file mode 100644 index 0000000000000..04c1b9ea7eba1 --- /dev/null +++ b/devel-common/src/tests_common/test_utils/otel_jaeger_utils.py @@ -0,0 +1,68 @@ +# 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. +"""Helpers for spans fetched from the Jaeger REST API.""" + +from __future__ import annotations + +from collections import defaultdict + + +def get_span_tags(span: dict) -> dict: + return {t["key"]: t["value"] for t in span.get("tags", [])} + + +def get_parent_span_id(span: dict) -> str | None: + """Return the spanID of the span's CHILD_OF parent, or None if it has no parent.""" + # Example: + # + # 'references': [{ + # 'refType': 'CHILD_OF', + # 'traceID': 'f94df3c57603264e8247844c6c95d115', + # 'spanID': '16737863fd381188' + # }] + for ref in span.get("references", []): + if ref["refType"] == "CHILD_OF": + return ref["spanID"] + return None + + +def get_descendant_span_names(parent_children_map: dict[str, list[dict]], span_id: str) -> set[str]: + """Recursively collect the operationNames of every span below span_id.""" + names = set() + for child in parent_children_map.get(span_id, []): + names.add(child["operationName"]) + names |= get_descendant_span_names(parent_children_map, child["spanID"]) + return names + + +def provided_child_spans_found_under_span(trace: dict, span_name: str, child_span_names: list[str]) -> bool: + """Return whether every child span appears in a Jaeger trace as a descendant of a given span_name.""" + # Dictionary of every span id, mapped to its direct children spans for the provided trace. + parent_children_map: dict[str, list[dict]] = defaultdict(list) + parent_span_id = None + for span in trace["spans"]: + span_parent_id = get_parent_span_id(span) + if span_parent_id is not None: + parent_children_map[span_parent_id].append(span) + if span["operationName"] == span_name: + parent_span_id = span["spanID"] + + if parent_span_id is None: + return False + + descendant_names = get_descendant_span_names(parent_children_map, parent_span_id) + return all(name in descendant_names for name in child_span_names) diff --git a/shared/observability/src/airflow_shared/observability/traces/__init__.py b/shared/observability/src/airflow_shared/observability/traces/__init__.py index ea866997db902..62895e0b53c36 100644 --- a/shared/observability/src/airflow_shared/observability/traces/__init__.py +++ b/shared/observability/src/airflow_shared/observability/traces/__init__.py @@ -29,14 +29,33 @@ from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter from opentelemetry.sdk.trace.id_generator import RandomIdGenerator from opentelemetry.sdk.trace.sampling import Decision -from opentelemetry.trace import NonRecordingSpan, Span, SpanContext, TraceFlags, TraceState +from opentelemetry.trace import INVALID_SPAN, NonRecordingSpan, Span, SpanContext, TraceFlags, TraceState from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator if TYPE_CHECKING: + from collections.abc import Iterator from configparser import ConfigParser log = logging.getLogger(__name__) tracer = trace.get_tracer(__name__) +_otel_debug_traces_on = False + + +@contextmanager +def start_debug_span(name: str, **kwargs) -> Iterator[Span]: + """ + Start a current span if the ``[traces] otel_debug_traces_on`` flag is enabled. + + Usable as a context manager or as a decorator. + When the flag is disabled, a non-recording span is yielded. + """ + if not _otel_debug_traces_on: + yield INVALID_SPAN + return + with tracer.start_as_current_span(name, **kwargs) as span: + yield span + + OVERRIDE_SPAN_ID_KEY = context.create_key("override_span_id") OVERRIDE_TRACE_ID_KEY = context.create_key("override_trace_id") @@ -249,10 +268,14 @@ def _load_exporter_from_env() -> SpanExporter: def configure_otel(conf: ConfigParser): + global _otel_debug_traces_on + otel_on = conf.getboolean("traces", "otel_on", fallback=False) if not otel_on: return + _otel_debug_traces_on = conf.getboolean("traces", "otel_debug_traces_on", fallback=False) + # ideally both endpoint and resource are None here # they would only be something other than None if user is using deprecated # Airflow-defined otel configs diff --git a/shared/observability/tests/observability/test_otel_utils.py b/shared/observability/tests/observability/test_otel_console_utils.py similarity index 99% rename from shared/observability/tests/observability/test_otel_utils.py rename to shared/observability/tests/observability/test_otel_console_utils.py index 556b7242a585a..f19f2b3d0742e 100644 --- a/shared/observability/tests/observability/test_otel_utils.py +++ b/shared/observability/tests/observability/test_otel_console_utils.py @@ -16,7 +16,7 @@ # under the License. from __future__ import annotations -from tests_common.test_utils.otel_utils import ( +from tests_common.test_utils.otel_console_utils import ( clean_task_lines, extract_metrics_from_output, extract_spans_from_output, @@ -26,7 +26,7 @@ ) -class TestUtilsUnit: +class TestConsoleUtilsUnit: # The method that extracts the spans from the output, # counts that there is no indentation on the cli, when a span starts and finishes. example_output = """ diff --git a/shared/observability/tests/observability/test_otel_jaeger_utils.py b/shared/observability/tests/observability/test_otel_jaeger_utils.py new file mode 100644 index 0000000000000..42eeedfbf45bd --- /dev/null +++ b/shared/observability/tests/observability/test_otel_jaeger_utils.py @@ -0,0 +1,173 @@ +# 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 pytest + +from tests_common.test_utils.otel_jaeger_utils import ( + get_descendant_span_names, + get_parent_span_id, + get_span_tags, + provided_child_spans_found_under_span, +) + +# A Jaeger trace with the following span hierarchy: +# +# scheduler.scheduler_loop +# ├── scheduler._do_scheduling +# │ └── scheduler.critical_section +# └── scheduler._process_executor_events +LOOP_SPAN = { + "spanID": "loop", + "operationName": "scheduler.scheduler_loop", + "references": [], + "tags": [ + {"key": "airflow.scheduler.loop_iteration.idle", "type": "bool", "value": False}, + {"key": "airflow.category", "type": "string", "value": "scheduler"}, + ], +} +DO_SCHEDULING_SPAN = { + "spanID": "do_scheduling", + "operationName": "scheduler._do_scheduling", + "references": [{"refType": "CHILD_OF", "traceID": "trace-1", "spanID": "loop"}], + "tags": [], +} +CRITICAL_SECTION_SPAN = { + "spanID": "critical_section", + "operationName": "scheduler.critical_section", + "references": [{"refType": "CHILD_OF", "traceID": "trace-1", "spanID": "do_scheduling"}], + "tags": [], +} +PROCESS_EVENTS_SPAN = { + "spanID": "process_events", + "operationName": "scheduler._process_executor_events", + "references": [{"refType": "CHILD_OF", "traceID": "trace-1", "spanID": "loop"}], + "tags": [], +} + +TRACE = { + "traceID": "trace-1", + "spans": [LOOP_SPAN, DO_SCHEDULING_SPAN, CRITICAL_SECTION_SPAN, PROCESS_EVENTS_SPAN], +} + +PARENT_CHILDREN_MAP = { + "loop": [DO_SCHEDULING_SPAN, PROCESS_EVENTS_SPAN], + "do_scheduling": [CRITICAL_SECTION_SPAN], +} + + +class TestJaegerUtilsUnit: + def test_get_span_tags(self): + assert get_span_tags(LOOP_SPAN) == { + "airflow.scheduler.loop_iteration.idle": False, + "airflow.category": "scheduler", + } + + @pytest.mark.parametrize( + "span", + [ + pytest.param({"tags": []}, id="empty-tags"), + pytest.param({}, id="missing-tags-key"), + ], + ) + def test_get_span_tags_without_tags(self, span: dict): + assert get_span_tags(span) == {} + + @pytest.mark.parametrize( + ("span", "expected_parent_id"), + [ + pytest.param(DO_SCHEDULING_SPAN, "loop", id="direct-child"), + pytest.param(CRITICAL_SECTION_SPAN, "do_scheduling", id="child-of-intermediate-span"), + pytest.param(LOOP_SPAN, None, id="root-span-empty-references"), + pytest.param({"operationName": "test"}, None, id="missing-references-key"), + pytest.param( + {"references": [{"refType": "FOLLOWS_FROM", "spanID": "other"}]}, + None, + id="non-child-of-ref-type", + ), + ], + ) + def test_get_parent_span_id(self, span: dict, expected_parent_id: str): + assert get_parent_span_id(span) == expected_parent_id + + @pytest.mark.parametrize( + ("span_id", "expected_names"), + [ + pytest.param( + "loop", + { + "scheduler._do_scheduling", + "scheduler.critical_section", + "scheduler._process_executor_events", + }, + id="root-collects-all-even-nested", + ), + pytest.param("do_scheduling", {"scheduler.critical_section"}, id="intermediate-node"), + pytest.param("critical_section", set(), id="leaf-has-no-descendants"), + pytest.param("unknown", set(), id="unknown-span-id"), + ], + ) + def test_get_descendant_span_names(self, span_id: str, expected_names: set): + assert get_descendant_span_names(PARENT_CHILDREN_MAP, span_id) == expected_names + + @pytest.mark.parametrize( + ("span_name", "child_span_names", "expected_bool_result"), + [ + pytest.param( + "scheduler.scheduler_loop", + ["scheduler._do_scheduling", "scheduler._process_executor_events"], + True, + id="all-direct-children-present", + ), + pytest.param( + "scheduler.scheduler_loop", + ["scheduler.critical_section"], + True, + id="nested-descendant-present", + ), + pytest.param( + "scheduler.scheduler_loop", + ["scheduler._do_scheduling", "scheduler.unknown_span"], + False, + id="unknown-child-span", + ), + pytest.param( + "scheduler.unknown_span", + ["scheduler._do_scheduling"], + False, + id="unknown-parent-span", + ), + pytest.param( + "scheduler._do_scheduling", + ["scheduler.critical_section"], + True, + id="descendant-under-intermediate-span", + ), + pytest.param( + "scheduler._do_scheduling", + ["scheduler._process_executor_events"], + False, + id="sibling-is-not-a-descendant", + ), + ], + ) + def test_provided_child_spans_found_under_span( + self, span_name: str, child_span_names: list, expected_bool_result + ): + assert ( + provided_child_spans_found_under_span(TRACE, span_name, child_span_names) is expected_bool_result + ) diff --git a/shared/observability/tests/observability/test_traces.py b/shared/observability/tests/observability/test_traces.py index eb0f499a250e1..b54eb241372b7 100644 --- a/shared/observability/tests/observability/test_traces.py +++ b/shared/observability/tests/observability/test_traces.py @@ -20,21 +20,25 @@ import pytest from opentelemetry import context, trace from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.sdk.trace.sampling import ( ALWAYS_OFF, ALWAYS_ON, ParentBased, TraceIdRatioBased, ) -from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags, TraceState +from opentelemetry.trace import INVALID_SPAN, NonRecordingSpan, SpanContext, TraceFlags, TraceState from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator +from airflow_shared.observability import traces from airflow_shared.observability.traces import ( DEFAULT_TASK_SPAN_DETAIL_LEVEL, TASK_SPAN_DETAIL_LEVEL_KEY, build_trace_state_entries, get_task_span_detail_level, new_dagrun_trace_carrier, + start_debug_span, ) @@ -360,3 +364,86 @@ def test_roundtrip_via_carrier(self): span = trace.get_current_span(ctx) assert get_task_span_detail_level(span) == 3 + + +@pytest.fixture +def set_debug_flag(monkeypatch): + """Set the module-level debug flag; monkeypatch restores it on teardown.""" + + def _set(value: bool): + monkeypatch.setattr(traces, "_otel_debug_traces_on", value) + + return _set + + +@pytest.fixture +def in_memory_exporter(monkeypatch): + """Use an in memory exporter and patch the module tracer to use it.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + monkeypatch.setattr(traces, "tracer", provider.get_tracer(__name__)) + return exporter + + +class TestStartDebugSpan: + @pytest.mark.parametrize( + ("debug_flag_val", "expected_span_num"), + [ + pytest.param(True, 1, id="flag_enabled"), + pytest.param(False, 0, id="flag_disabled"), + ], + ) + def test_start_debug_span_context_manager( + self, debug_flag_val: bool, expected_span_num: int, set_debug_flag, in_memory_exporter + ): + set_debug_flag(debug_flag_val) + span_name = "ctx_mgr_test_span" + + # If the flag is enabled, then the span will be recording, else it won't be. + is_recording = debug_flag_val + + with start_debug_span("ctx_mgr_test_span") as span: + if expected_span_num > 0: + expected_span = trace.get_current_span() + else: + expected_span = INVALID_SPAN + assert span.is_recording() is is_recording + assert span is expected_span + # If the span is non-recording, enriching it is a safe no-op + span.set_attribute("a", 1) + + finished_spans = in_memory_exporter.get_finished_spans() + + assert len(finished_spans) == expected_span_num + if expected_span_num > 0: + assert finished_spans[0].name == span_name + assert finished_spans[0].attributes["a"] == 1 + + @pytest.mark.parametrize( + ("debug_flag_val", "expected_span_num"), + [ + pytest.param(True, 1, id="flag_enabled"), + pytest.param(False, 0, id="flag_disabled"), + ], + ) + def test_start_debug_span_decorator( + self, debug_flag_val: bool, expected_span_num: int, set_debug_flag, in_memory_exporter + ): + span_name = "decorator_test_span" + + @start_debug_span(span_name) + def double(x): + trace.get_current_span().set_attribute("a", 1) + return x * 2 + + set_debug_flag(debug_flag_val) + + assert double(2) == 4 + + finished_spans = in_memory_exporter.get_finished_spans() + + assert len(finished_spans) == expected_span_num + if expected_span_num > 0: + assert finished_spans[0].name == span_name + assert finished_spans[0].attributes["a"] == 1