diff --git a/airflow-core/src/airflow/executors/workloads/task.py b/airflow-core/src/airflow/executors/workloads/task.py index 68a917118e39c..3099fe1d77485 100644 --- a/airflow-core/src/airflow/executors/workloads/task.py +++ b/airflow-core/src/airflow/executors/workloads/task.py @@ -107,7 +107,11 @@ def make( bundle_info = BundleInfo( name=ti.dag_model.bundle_name, version=ti.dag_run.bundle_version, - version_data=_resolve_version_data(ti.dag_version, ti.dag_run.bundle_version), + # Source version_data from the run's pinned version (matching ``version`` above), + # not the TI's dag_version. A mid-run DAG re-parse can bump the TI's dag_version + # to a newer version while the run stays pinned; sourcing from created_dag_version + # keeps the shipped hash and manifest consistent so versioned bundles stay reproducible. + version_data=_resolve_version_data(ti.dag_run.created_dag_version, ti.dag_run.bundle_version), ) fname = log_filename_template_renderer()(ti=ti) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index f6d6a4787f629..c2f3ee05d730e 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -684,12 +684,17 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - ranked_query.c.map_index_for_ordering, ) .options(selectinload(TI.dag_model)) - # Eager-load dag_version: TIs become transient (via make_transient) before - # ExecuteTask.make() reads ti.dag_version.version_data. Lazy loads on - # transient objects silently return None instead of raising DetachedInstanceError. - # Scope the second SELECT to version_data (the PK is auto-included) so we read - # two columns rather than the full DagVersion row. - .options(selectinload(TI.dag_version).load_only(DagVersion.version_data)) + # Eager-load the run's pinned DagVersion (dag_run.created_dag_version): TIs become + # transient (via make_transient) before ExecuteTask.make() reads + # ti.dag_run.created_dag_version.version_data to ship the bundle manifest matching + # the run's pinned bundle_version. Lazy loads on transient objects silently return + # None instead of raising DetachedInstanceError. Scope the SELECT to version_data + # (the PK is auto-included) so we read two columns rather than the full row. + .options( + joinedload(TI.dag_run) + .selectinload(DagRun.created_dag_version) + .load_only(DagVersion.version_data) + ) ) query = query.limit(max_tis) diff --git a/airflow-core/tests/unit/executors/test_workloads.py b/airflow-core/tests/unit/executors/test_workloads.py index 1ef04477d081c..37fbcd96ce950 100644 --- a/airflow-core/tests/unit/executors/test_workloads.py +++ b/airflow-core/tests/unit/executors/test_workloads.py @@ -184,12 +184,21 @@ def _stub_log_template(self, monkeypatch): ) @staticmethod - def _make_mock_ti(bundle_version, version_data, *, has_dag_version=True): + def _make_mock_ti( + bundle_version, + version_data, + *, + has_created_dag_version=True, + ti_dag_version_data=None, + ): """Build a mock TI with the attributes ExecuteTask.make() reads. - ``has_dag_version`` controls whether the TI has an associated DagVersion - (legacy/backfilled TIs may not), independently of ``version_data`` so the - pin-guard can be exercised with version_data present on an unpinned run. + ``version_data`` is the manifest on the run's pinned version + (``dag_run.created_dag_version``) -- the source make() must use. + ``ti_dag_version_data`` is the manifest on ``ti.dag_version``, which make() + must IGNORE (it can diverge from the run's pin after a mid-run DAG re-parse). + ``has_created_dag_version`` toggles whether the run has a pinned DagVersion + (legacy/backfilled runs may not). """ from unittest.mock import Mock @@ -215,15 +224,18 @@ def _make_mock_ti(bundle_version, version_data, *, has_dag_version=True): ti.dag_run.bundle_version = bundle_version - if has_dag_version: - ti.dag_version.version_data = version_data + # make() must source version_data from the run's pinned version, never ti.dag_version. + ti.dag_version.version_data = ti_dag_version_data + + if has_created_dag_version: + ti.dag_run.created_dag_version.version_data = version_data else: - ti.dag_version = None + ti.dag_run.created_dag_version = None return ti def test_pinned_run_populates_version_data(self): - """When the run is pinned, version_data from dag_version flows to BundleInfo.""" + """When the run is pinned, version_data from the run's created_dag_version flows to BundleInfo.""" version_data = {"schema_version": 1, "files": {"dags/my_dag.py": "ver123"}} ti = self._make_mock_ti(bundle_version="abc123", version_data=version_data) @@ -233,7 +245,7 @@ def test_pinned_run_populates_version_data(self): assert workload.bundle_info.version_data == version_data def test_unpinned_run_suppresses_present_version_data(self): - """An unpinned run must not expose version_data even when the dag_version carries it.""" + """An unpinned run must not expose version_data even when created_dag_version carries it.""" version_data = {"schema_version": 1, "files": {"dags/my_dag.py": "ver123"}} ti = self._make_mock_ti(bundle_version=None, version_data=version_data) @@ -242,15 +254,34 @@ def test_unpinned_run_suppresses_present_version_data(self): assert workload.bundle_info.version is None assert workload.bundle_info.version_data is None - def test_missing_dag_version_yields_none(self): - """A pinned run whose TI has no dag_version (legacy/backfilled) yields no version_data.""" - ti = self._make_mock_ti(bundle_version="abc123", version_data=None, has_dag_version=False) + def test_missing_created_dag_version_yields_none(self): + """A pinned run whose DagRun has no created_dag_version yields no version_data.""" + ti = self._make_mock_ti(bundle_version="abc123", version_data=None, has_created_dag_version=False) workload = ExecuteTask.make(ti) assert workload.bundle_info.version == "abc123" assert workload.bundle_info.version_data is None + def test_mid_run_dag_version_bump_uses_run_pinned_manifest(self): + """Regression: a mid-run DAG re-parse can bump ti.dag_version to a newer version while the + run stays pinned. make() must ship the run's pinned manifest (created_dag_version), not the + TI's bumped one -- otherwise a versioned bundle would fetch the wrong (latest) code. + """ + run_manifest = {"schema_version": 1, "files": {"dags/my_dag.py": "v1-object-id"}} + bumped_manifest = {"schema_version": 1, "files": {"dags/my_dag.py": "v2-object-id"}} + ti = self._make_mock_ti( + bundle_version="v1hash", + version_data=run_manifest, + ti_dag_version_data=bumped_manifest, + ) + + workload = ExecuteTask.make(ti) + + assert workload.bundle_info.version == "v1hash" + assert workload.bundle_info.version_data == run_manifest + assert workload.bundle_info.version_data != bumped_manifest + class TestExecuteCallbackMakeVersionData: """Tests for ExecuteCallback.make() threading version_data through BundleInfo.""" diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 1f3b6992c3e1c..68f4ec0915092 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -142,7 +142,7 @@ from airflow.utils.types import DagRunTriggeredByType, DagRunType from tests_common.pytest_plugin import AIRFLOW_ROOT_PATH -from tests_common.test_utils.asserts import assert_queries_count +from tests_common.test_utils.asserts import assert_queries_count, count_queries from tests_common.test_utils.config import conf_vars, env_vars from tests_common.test_utils.dag import create_scheduler_dag, sync_dag_to_db, sync_dags_to_db from tests_common.test_utils.db import ( @@ -1425,6 +1425,41 @@ def test_find_executable_task_instances_backfill(self, dag_maker): assert {x.key for x in queued_tis} == {ti_non_backfill.key, ti_backfill.key} session.rollback() + def test_executable_task_instances_no_per_ti_queries(self, dag_maker, session): + """Guard against an N+1 when enqueuing task instances. + + ``ExecuteTask.make()`` reads ``ti.dag_run.created_dag_version.version_data`` to ship the + run's pinned bundle manifest. ``dag_run`` is eager-joined and ``created_dag_version`` is a + single batched ``selectin``, so the number of queries in + ``_executable_task_instances_to_queued`` must be independent of how many task instances are + in the batch. If a future change lazy-loads ``dag_run``/``created_dag_version`` per TI, the + count would scale with the task count and this test fails. + """ + scheduler_job = Job() + runner = SchedulerJobRunner(job=scheduler_job) + self.job_runner = runner + + def _measure(dag_id: str, num_tasks: int) -> int: + with dag_maker(dag_id=dag_id, max_active_tasks=64, session=session): + for i in range(num_tasks): + EmptyOperator(task_id=f"t{i}") + dr = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) + for ti in dr.task_instances: + ti.state = State.SCHEDULED + session.flush() + with count_queries(session=session) as result: + runner._executable_task_instances_to_queued(max_tis=64, session=session) + session.rollback() + return sum(result.values()) + + one_task = _measure("q_count_one", 1) + many_tasks = _measure("q_count_many", 10) + + assert one_task == many_tasks, ( + f"query count scaled with task-instance count ({one_task} -> {many_tasks}); " + "likely a per-TI lazy load (N+1) of dag_run/created_dag_version" + ) + def test_find_executable_task_instances_mysql_hint_only_applies_to_inner_query(self, dag_maker, session): dag_id = "SchedulerJobTest.test_find_executable_task_instances_mysql_hint_only_applies_to_inner_query" task_id = "dummy" diff --git a/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py b/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py index f5800232610ca..4bfd7c716efa0 100644 --- a/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py +++ b/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py @@ -817,6 +817,8 @@ def test_try_adopt_task_instances(self, mock_executor): task.dag_version = mock.Mock(version_data=None) task.dag_run = mock.Mock() task.dag_run.bundle_version = "1.0.0" + # ExecuteTask.make() sources version_data from the run's pinned version. + task.dag_run.created_dag_version = mock.Mock(version_data=None) task.dag_run.context_carrier = {} if not AIRFLOW_V_3_0_PLUS: diff --git a/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py b/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py index b7c25d51277b4..b6074b8a55402 100644 --- a/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py +++ b/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py @@ -1312,6 +1312,8 @@ def test_try_adopt_task_instances(self, mock_executor): task.dag_version = mock.Mock(version_data=None) task.dag_run = mock.Mock() task.dag_run.bundle_version = "1.0.0" + # ExecuteTask.make() sources version_data from the run's pinned version. + task.dag_run.created_dag_version = mock.Mock(version_data=None) task.dag_run.context_carrier = {} # Mock command generation based on Airflow version