This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new fb69e2b4a32 Pin task bundle manifest to the dagrun's version (#69941)
fb69e2b4a32 is described below
commit fb69e2b4a322f790504f049622b8cd9eb316bdcf
Author: Niko Oliveira <[email protected]>
AuthorDate: Fri Jul 17 12:35:30 2026 -0700
Pin task bundle manifest to the dagrun's version (#69941)
A mid-run DAG re-parse possibly creates a new DagVersion and bumps
not-yet-started task instances' dag_version_id, while the DagRun
stays pinned to its original bundle_version. ExecuteTask.make()
sourced the bundle version from the run's pin (dag_run.bundle_version)
but bundle version_data from the task's now-bumped dag_version,
so the workload shipped a hash and a manifest that describe
different versions. This double source was noticed for the callback
path in a previous PR, but the task path was missed.
Source version_data from the run's pinned created_dag_version so it always
matches the bundle version, mirroring ExecuteCallback.make() and the way Git
bundles resolve everything from the run's pinned commit.
---
.../src/airflow/executors/workloads/task.py | 6 ++-
.../src/airflow/jobs/scheduler_job_runner.py | 17 ++++---
.../tests/unit/executors/test_workloads.py | 55 +++++++++++++++++-----
airflow-core/tests/unit/jobs/test_scheduler_job.py | 37 ++++++++++++++-
.../aws/executors/batch/test_batch_executor.py | 2 +
.../amazon/aws/executors/ecs/test_ecs_executor.py | 2 +
6 files changed, 99 insertions(+), 20 deletions(-)
diff --git a/airflow-core/src/airflow/executors/workloads/task.py
b/airflow-core/src/airflow/executors/workloads/task.py
index 68a917118e3..3099fe1d774 100644
--- a/airflow-core/src/airflow/executors/workloads/task.py
+++ b/airflow-core/src/airflow/executors/workloads/task.py
@@ -107,7 +107,11 @@ class ExecuteTask(BaseDagBundleWorkload):
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 f6d6a4787f6..c2f3ee05d73 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 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
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 1ef04477d08..37fbcd96ce9 100644
--- a/airflow-core/tests/unit/executors/test_workloads.py
+++ b/airflow-core/tests/unit/executors/test_workloads.py
@@ -184,12 +184,21 @@ class TestExecuteTaskMakeVersionData:
)
@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 @@ class TestExecuteTaskMakeVersionData:
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 @@ class TestExecuteTaskMakeVersionData:
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 @@ class TestExecuteTaskMakeVersionData:
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 1f3b6992c3e..68f4ec09150 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.state import CallbackState, DagRunState,
State, TaskInstanceS
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 @@ class TestSchedulerJob:
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 f5800232610..4bfd7c716ef 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 @@ class TestAwsBatchExecutor:
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 b7c25d51277..b6074b8a554 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 @@ class TestAwsEcsExecutor:
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
