This is an automated email from the ASF dual-hosted git repository.
bbovenzi 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 05981d3e2aa Fix task Clear and callback bundles after code-only Dag
changes (#73915)
05981d3e2aa is described below
commit 05981d3e2aaa2a4ac4aa5c9a5fba350459b5803b
Author: Dheeraj Turaga <[email protected]>
AuthorDate: Tue Oct 20 18:04:29 2026 -0500
Fix task Clear and callback bundles after code-only Dag changes (#73915)
* Fix task Clear bundle selection after code-only Dag changes
A deployed callable fix can leave serialization unchanged while an existing
run stays pinned to an older bundle. Users need the latest-version clear
option to remain available and an explicit selection to take effect for
active runs. Scheduler callbacks also need to resolve the same bundle as
task execution.
* Pair task callback bundle manifest with the run's pinned version
Callbacks and failure emails resolve their bundle version from the Dag run,
but the manifest still came from the task instance's Dag version. When that
version differs from the one the run was created with, a manifest-based
bundle would initialize the run's bundle version from another version's
manifest. Task execution and Dag-run callbacks already take both from the
run's pinned version.
* Run triggerer task callbacks on the run's pinned bundle
A deferred task that finishes in the triggerer has its callback sent from
there, and that path still took the bundle from the task instance's Dag
version. A code-only update advances that version's bundle in place while
the run stays pinned, so the callback would run newer code than the task
did. Scheduler callbacks already resolve the bundle from the run.
* Stop pairing a pinned bundle with an overwritten manifest
A code-only update reuses the Dag version and overwrites its bundle version
and manifest in place, while existing runs stay pinned to the older bundle
version. Task execution and callbacks then shipped the run's older version
with the newer manifest, so a manifest-based bundle would materialize code
the run was never pinned to. Once the original manifest is gone, sending
none is safer than sending one for a different version.
---
.../src/airflow/jobs/scheduler_job_runner.py | 60 ++-------
airflow-core/src/airflow/models/dag_version.py | 22 +++-
airflow-core/src/airflow/models/serialized_dag.py | 6 +-
airflow-core/src/airflow/models/trigger.py | 21 +---
.../TaskInstance/ClearTaskInstanceDialog.test.tsx | 137 ++++++++++++++++++---
.../Clear/TaskInstance/ClearTaskInstanceDialog.tsx | 9 +-
.../tests/unit/executors/test_workloads.py | 2 +
airflow-core/tests/unit/jobs/test_scheduler_job.py | 92 +++++++++++---
airflow-core/tests/unit/models/test_dag_version.py | 8 +-
airflow-core/tests/unit/models/test_trigger.py | 63 +++++++++-
10 files changed, 308 insertions(+), 112 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 7729ca789e5..7449c09e348 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -105,7 +105,7 @@ from airflow.models.connection_test import (
ConnectionTestState,
)
from airflow.models.dag import DagModel
-from airflow.models.dag_version import DagVersion, _resolve_version_data
+from airflow.models.dag_version import DagVersion,
_resolve_ti_callback_bundle_info
from airflow.models.dagbag import CachedDBDagBag, DBDagBag
from airflow.models.dagbundle import DagBundleModel
from airflow.models.dagrun import DagRun
@@ -200,26 +200,6 @@ def _eager_load_dag_run_for_validation() ->
tuple[LoaderOption, LoaderOption]:
)
-def _resolve_ti_callback_bundle_info(ti: TaskInstance) -> tuple[str, str |
None, Any]:
- """
- Resolve the bundle name/version/version-data needed to build a
TaskCallbackRequest or EmailRequest.
-
- Used by the heartbeat-timeout purge path. Encapsulates the bundle-pinning
semantics: fall back
- to ``dag_model`` for legacy tasks with no ``dag_version`` (pre-AIP-66
migrations), and leave the
- bundle version unpinned when the dag run itself wasn't pinned
(``disable_bundle_versioning``),
- so the callback runs against the same code as the task did.
``process_executor_events`` inlines
- the same resolution for its externally-killed-task path.
- """
- bundle_name = ti.dag_version.bundle_name if ti.dag_version else
ti.dag_model.bundle_name
- bundle_version = (
- ti.dag_version.bundle_version
- if ti.dag_version and ti.dag_run.bundle_version is not None
- else ti.dag_run.bundle_version
- )
- version_data = _resolve_version_data(ti.dag_version,
ti.dag_run.bundle_version)
- return bundle_name, bundle_version, version_data
-
-
def _ensure_ti_has_dag_version_id(ti: TaskInstance, session: Session, log:
Logger) -> bool:
"""
Ensure a TaskInstance has a valid dag_version_id for Pydantic
serialisation.
@@ -1733,18 +1713,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
# Only log the error/extra info here, since the
`ti.handle_failure()` path will log it
# too, which would lead to double logging
cls.logger().error(msg)
- # Safely extract bundle info: prefer dag_version when
available,
- # fall back to dag_model/dag_run for legacy tasks migrated
from
- # Airflow 2 where dag_version may be None (AIP-66).
- _bundle_name = ti.dag_version.bundle_name if
ti.dag_version else ti.dag_model.bundle_name
- # Mirror dag_run pinning: if the run wasn't pinned (e.g.
dag.disable_bundle_versioning=True),
- # leave the callback unpinned so it runs against the same
code as the task.
- _bundle_version = (
- ti.dag_version.bundle_version
- if ti.dag_version and ti.dag_run.bundle_version is not
None
- else ti.dag_run.bundle_version
- )
- _version_data = _resolve_version_data(ti.dag_version,
ti.dag_run.bundle_version)
+ _bundle_name, _bundle_version, _version_data =
_resolve_ti_callback_bundle_info(ti)
# Backfill dag_version_id for legacy tasks (Pydantic
requires uuid.UUID).
if not _ensure_ti_has_dag_version_id(ti, session,
cls.logger()):
continue
@@ -1778,15 +1747,9 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
"Sending email request for task %s to DAG Processor",
ti,
)
- # Safely extract bundle info with fallback for legacy tasks
- # (dag_version may be None after Airflow 2 → 3 migration).
- _email_bundle_name = (
- ti.dag_version.bundle_name if ti.dag_version else
ti.dag_model.bundle_name
- )
- _email_bundle_version = (
- ti.dag_version.bundle_version if ti.dag_version else
ti.dag_run.bundle_version
+ _email_bundle_name, _email_bundle_version,
_email_version_data = (
+ _resolve_ti_callback_bundle_info(ti)
)
- _email_version_data =
_resolve_version_data(ti.dag_version, ti.dag_run.bundle_version)
# Backfill dag_version_id for legacy tasks (Pydantic
requires uuid.UUID).
if not _ensure_ti_has_dag_version_id(ti, session,
cls.logger()):
continue
@@ -3496,19 +3459,9 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
if task.has_on_failure_callback:
if inspect(ti).detached:
ti = session.merge(ti)
- # Safely extract bundle info with fallback for legacy tasks
- # (dag_version may be None after Airflow 2 → 3 migration).
- _stuck_bundle_name = (
- ti.dag_version.bundle_name if ti.dag_version else
ti.dag_model.bundle_name
+ _stuck_bundle_name, _stuck_bundle_version,
_stuck_version_data = (
+ _resolve_ti_callback_bundle_info(ti)
)
- # Mirror dag_run pinning: if the run wasn't pinned (e.g.
dag.disable_bundle_versioning=True),
- # leave the callback unpinned so it runs against the same
code as the task.
- _stuck_bundle_version = (
- ti.dag_version.bundle_version
- if ti.dag_version and ti.dag_run.bundle_version is not
None
- else ti.dag_run.bundle_version
- )
- _stuck_version_data =
_resolve_version_data(ti.dag_version, ti.dag_run.bundle_version)
# Backfill dag_version_id for legacy tasks (Pydantic
requires uuid.UUID).
# Note: we cannot use `continue` here because this method
is not
# inside a loop. If backfilling fails we simply skip the
callback.
@@ -3980,6 +3933,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
.options(selectinload(TI.dag_model))
.options(asset_loader)
.options(alias_loader)
+
.options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
.options(selectinload(TI.dag_version))
.with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
.join(DM, TI.dag_id == DM.dag_id)
diff --git a/airflow-core/src/airflow/models/dag_version.py
b/airflow-core/src/airflow/models/dag_version.py
index ff4971f4ba8..8dc6376a185 100644
--- a/airflow-core/src/airflow/models/dag_version.py
+++ b/airflow-core/src/airflow/models/dag_version.py
@@ -41,6 +41,7 @@ if TYPE_CHECKING:
from sqlalchemy.sql import Select
from airflow.models.serialized_dag import SerializedDagModel
+ from airflow.models.taskinstance import TaskInstance
log = logging.getLogger(__name__)
@@ -416,7 +417,24 @@ def _resolve_version_data(
"""Return a bundle version's ``version_data`` manifest, but only for
pinned runs."""
# Expose version_data only when the run is pinned (bundle_version set) and
a DagVersion is
# present, so the bundle initializes against the exact version the run
used. Unpinned runs
- # follow the latest bundle state, and legacy rows have no DagVersion.
- if dag_version is not None and bundle_version is not None:
+ # follow the latest bundle state, and legacy rows have no DagVersion. A
code-only update
+ # overwrites a reused DagVersion's bundle metadata in place, so its
manifest describes the
+ # pinned bundle only while the two bundle versions still match.
+ if (
+ dag_version is not None
+ and bundle_version is not None
+ and dag_version.bundle_version == bundle_version
+ ):
return dag_version.version_data
return None
+
+
+def _resolve_ti_callback_bundle_info(
+ ti: TaskInstance,
+) -> tuple[str, str | None, dict[str, Any] | None]:
+ """Resolve bundle metadata for task callbacks and emails using the run's
execution bundle."""
+ bundle_name = ti.dag_version.bundle_name if ti.dag_version else
ti.dag_model.bundle_name
+ # A code-only update can advance the DagVersion's bundle while the run
stays pinned.
+ bundle_version = ti.dag_run.bundle_version
+ version_data = _resolve_version_data(ti.dag_run.created_dag_version,
ti.dag_run.bundle_version)
+ return bundle_name, bundle_version, version_data
diff --git a/airflow-core/src/airflow/models/serialized_dag.py
b/airflow-core/src/airflow/models/serialized_dag.py
index aa2e801d7d5..6cab8defb81 100644
--- a/airflow-core/src/airflow/models/serialized_dag.py
+++ b/airflow-core/src/airflow/models/serialized_dag.py
@@ -705,10 +705,8 @@ class SerializedDagModel(Base):
new_dag_hash = cls.hash(dag.data)
if serialized_dag_hash == new_dag_hash and dag_version and
dag_version.bundle_name == bundle_name:
- # Serialized content is unchanged, so we don't create a new
DagVersion.
- # But if the bundle advanced, refresh the latest version's pointer
in place — tasks resolve
- # their code from ``ti.dag_version.bundle_version`` at run time,
so a stale
- # pointer makes runs execute an outdated commit.
+ # Reuse the structural version with the latest bundle metadata.
Existing runs
+ # keep their execution bundle pinned separately in
DagRun.bundle_version.
bundle_metadata_changed = (
dag_version.bundle_version != bundle_version or
dag_version.version_data != version_data
)
diff --git a/airflow-core/src/airflow/models/trigger.py
b/airflow-core/src/airflow/models/trigger.py
index a7193074f6f..656a2a474a1 100644
--- a/airflow-core/src/airflow/models/trigger.py
+++ b/airflow-core/src/airflow/models/trigger.py
@@ -722,24 +722,9 @@ def _(event: BaseTaskEndEvent, *, task_instance:
TaskInstance, session: Session)
):
if task_instance.dag_model.relative_fileloc is None:
raise RuntimeError("relative_fileloc should not be None for a
finished task")
- from airflow.models.dag_version import _resolve_version_data
-
- # Derive bundle identity from the TI's dag_version (falling back
to dag_run/dag_model
- # for legacy/unpinned runs), mirroring the other callback sites so
bundle_version and
- # version_data always describe the same version.
- bundle_name = (
- task_instance.dag_version.bundle_name
- if task_instance.dag_version
- else task_instance.dag_model.bundle_name
- )
- bundle_version = (
- task_instance.dag_version.bundle_version
- if task_instance.dag_version and
task_instance.dag_run.bundle_version is not None
- else task_instance.dag_run.bundle_version
- )
- version_data = _resolve_version_data(
- task_instance.dag_version, task_instance.dag_run.bundle_version
- )
+ from airflow.models.dag_version import
_resolve_ti_callback_bundle_info
+
+ bundle_name, bundle_version, version_data =
_resolve_ti_callback_bundle_info(task_instance)
request = TaskCallbackRequest(
filepath=task_instance.dag_model.relative_fileloc,
ti=task_instance_to_runtime(
diff --git
a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.test.tsx
b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.test.tsx
index 0d5135cedf5..b585bd6a5ec 100644
---
a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.test.tsx
+++
b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.test.tsx
@@ -21,7 +21,7 @@ import { fireEvent, render, screen, waitFor, within } from
"@testing-library/rea
import { http, HttpResponse } from "msw";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import type { TaskInstanceResponse } from "openapi/requests/types.gen";
+import type { DagVersionResponse, TaskInstanceResponse } from
"openapi/requests/types.gen";
import BulkClearTaskInstancesButton from
"src/pages/TaskInstances/BulkClearTaskInstancesButton";
@@ -93,6 +93,18 @@ const loopMember: TaskInstanceResponse = {
in_loop: true,
};
+// A code-only change advanced the Dag version's bundle in place; the run
stays pinned to its bundle.
+const latestVersion: DagVersionResponse = {
+ bundle_name: "git-bundle",
+ bundle_url: null,
+ bundle_version: "latest-sha",
+ created_at: "2025-01-02T00:00:00Z",
+ dag_display_name: "Test DAG",
+ dag_id: DAG_ID,
+ id: "version-21",
+ version_number: 21,
+};
+
const mutateMock = vi.fn();
const dryRunsMock = vi.fn();
const confirmationDryRun = { isFetching: false };
@@ -131,6 +143,29 @@ afterEach(() => {
localStorage.clear();
});
+const serveBundleOnlyUpdate = (runLoaded: Promise<void> = Promise.resolve()) =>
+ server.use(
+ http.get(`/api/v2/dags/${DAG_ID}/details`, () =>
+ HttpResponse.json({
+ bundle_version: "latest-sha",
+ latest_dag_version: latestVersion,
+ rerun_with_latest_version: null,
+ }),
+ ),
+ http.get(`/api/v2/dags/${DAG_ID}/dagRuns/${DAG_RUN_ID}`, async () => {
+ await runLoaded;
+
+ return HttpResponse.json({ bundle_version: "historical-sha",
dag_versions: [latestVersion] });
+ }),
+ );
+
+const clickConfirm = async () => {
+ const confirmButton = await screen.findByRole("button", { name:
/modal\.confirm/iu });
+
+ await waitFor(() => expect(confirmButton).toBeEnabled());
+ fireEvent.click(confirmButton);
+};
+
describe("ClearTaskInstanceDialog", () => {
it("seeds the keep-task-state checkbox from the stored default and sends it
on confirm", async () => {
localStorage.setItem(CLEAR_KEEP_TASK_STATE_KEY, JSON.stringify(true));
@@ -143,9 +178,7 @@ describe("ClearTaskInstanceDialog", () => {
expect(keepTaskStateCheckbox).toBeChecked();
- const confirmButton = await screen.findByRole("button", { name:
/modal\.confirm/iu });
-
- fireEvent.click(confirmButton);
+ await clickConfirm();
await waitFor(() => expect(mutateMock).toHaveBeenCalled());
@@ -165,7 +198,7 @@ describe("ClearTaskInstanceDialog", () => {
},
);
- fireEvent.click(await screen.findByRole("button", { name:
/modal\.confirm/iu }));
+ await clickConfirm();
// Let the clear dialog finish closing while the running-task check is
still in flight.
await new Promise((resolve) => {
setTimeout(resolve, 0);
@@ -186,7 +219,7 @@ describe("ClearTaskInstanceDialog", () => {
wrapper: Wrapper,
});
- fireEvent.click(await screen.findByRole("button", { name:
/modal\.confirm/iu }));
+ await clickConfirm();
expect(await
screen.findByText(/confirmationDialog\.title/iu)).toBeVisible();
expect(mutateMock).not.toHaveBeenCalled();
@@ -202,7 +235,7 @@ describe("ClearTaskInstanceDialog", () => {
const rows = await screen.findAllByRole("row");
fireEvent.click(within(rows[1] as HTMLElement).getByRole("checkbox"));
- fireEvent.click(await screen.findByRole("button", { name:
/modal\.confirm/iu }));
+ await clickConfirm();
await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1));
const [{ requestBody }] = mutateMock.mock.calls[0] as [
@@ -226,7 +259,7 @@ describe("ClearTaskInstanceDialog", () => {
affectedTasks.task_instances = [loopPass("pass-0-retried", 0),
loopPass("pass-2", 2)];
rerender(dialog);
- fireEvent.click(await screen.findByRole("button", { name:
/modal\.confirm/iu }));
+ await clickConfirm();
await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1));
const [{ requestBody }] = mutateMock.mock.calls[0] as [
@@ -280,7 +313,7 @@ describe("ClearTaskInstanceDialog", () => {
});
fireEvent.click(await screen.findByRole("button", { name:
/options\.laterIterations/iu }));
- fireEvent.click(await screen.findByRole("button", { name:
/modal\.confirm/iu }));
+ await clickConfirm();
await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1));
const [{ requestBody }] = mutateMock.mock.calls[0] as [{ requestBody:
Record<string, unknown> }];
@@ -304,7 +337,7 @@ describe("ClearTaskInstanceDialog", () => {
});
fireEvent.click(await screen.findByRole("button", { name:
/options\.wholeExpansion/iu }));
- fireEvent.click(await screen.findByRole("button", { name:
/modal\.confirm/iu }));
+ await clickConfirm();
await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1));
const [{ requestBody }] = mutateMock.mock.calls[0] as [{ requestBody:
Record<string, unknown> }];
@@ -344,7 +377,7 @@ describe("ClearTaskInstanceDialog", () => {
wrapper: Wrapper,
});
- fireEvent.click(await screen.findByRole("button", { name:
/modal\.confirm/iu }));
+ await clickConfirm();
await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1));
const [{ requestBody }] = mutateMock.mock.calls[0] as [{ requestBody:
Record<string, unknown> }];
@@ -366,7 +399,7 @@ describe("ClearTaskInstanceDialog", () => {
wrapper: Wrapper,
});
- fireEvent.click(await screen.findByRole("button", { name:
/modal\.confirm/iu }));
+ await clickConfirm();
await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(2));
expect(
@@ -393,7 +426,7 @@ describe("ClearTaskInstanceDialog", () => {
const note = await screen.findByRole("textbox", { hidden: true });
expect(note).toHaveValue("existing note");
- fireEvent.click(await screen.findByRole("button", { name:
/modal\.confirm/iu }));
+ await clickConfirm();
await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1));
expect((mutateMock.mock.calls[0] as [{ requestBody: { note?: string }
}])[0].requestBody.note).toBe(
undefined,
@@ -408,7 +441,7 @@ describe("ClearTaskInstanceDialog", () => {
fireEvent.change(await screen.findByRole("textbox", { hidden: true }), {
target: { value: "new reason" },
});
- fireEvent.click(await screen.findByRole("button", { name:
/modal\.confirm/iu }));
+ await clickConfirm();
await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1));
expect((mutateMock.mock.calls[0] as [{ requestBody: { note?: string }
}])[0].requestBody.note).toBe(
"new reason",
@@ -494,6 +527,82 @@ describe("ClearTaskInstanceDialog", () => {
},
);
+ it.each([false, true])(
+ "offers the run-on-latest option when only the run's bundle is outdated
and sends the choice: %s",
+ async (runOnLatest) => {
+ serveBundleOnlyUpdate();
+ render(
+ <ClearTaskInstanceDialog
+ onClose={vi.fn()}
+ open
+ taskInstance={{ ...taskInstance, dag_version: latestVersion }}
+ />,
+ { wrapper: Wrapper },
+ );
+
+ const checkbox = await screen.findByRole("checkbox", { name:
/runOnLatestVersion/iu });
+
+ expect(checkbox).not.toBeChecked();
+ if (runOnLatest) {
+ fireEvent.click(checkbox);
+ await waitFor(() => expect(checkbox).toBeChecked());
+ }
+ await clickConfirm();
+
+ await waitFor(() => expect(mutateMock).toHaveBeenCalledTimes(1));
+ expect((mutateMock.mock.calls[0] as [{ requestBody: unknown
}])[0].requestBody).toMatchObject({
+ dag_run_id: DAG_RUN_ID,
+ dry_run: false,
+ run_on_latest_version: runOnLatest,
+ });
+ },
+ );
+
+ it("offers the run-on-latest option when clearing all mapped instances on an
outdated bundle", async () => {
+ serveBundleOnlyUpdate();
+ render(
+ <ClearTaskInstanceDialog
+ allMapped
+ dagId={DAG_ID}
+ dagRunId={DAG_RUN_ID}
+ onClose={vi.fn()}
+ open
+ taskId={TASK_ID}
+ />,
+ { wrapper: Wrapper },
+ );
+
+ expect(await screen.findByRole("checkbox", { name: /runOnLatestVersion/iu
})).toBeVisible();
+ });
+
+ it("keeps confirm disabled until the run's bundle has loaded", async () => {
+ let loadRun: (() => void) | undefined;
+
+ serveBundleOnlyUpdate(
+ new Promise<void>((resolve) => {
+ loadRun = resolve;
+ }),
+ );
+ render(
+ <ClearTaskInstanceDialog
+ onClose={vi.fn()}
+ open
+ taskInstance={{ ...taskInstance, dag_version: latestVersion }}
+ />,
+ { wrapper: Wrapper },
+ );
+
+ const confirmButton = await screen.findByRole("button", { name:
/modal\.confirm/iu });
+
+ expect(confirmButton).toBeDisabled();
+ expect(screen.queryByRole("checkbox", { name: /runOnLatestVersion/iu
})).not.toBeInTheDocument();
+
+ loadRun?.();
+
+ expect(await screen.findByRole("checkbox", { name: /runOnLatestVersion/iu
})).toBeVisible();
+ await waitFor(() => expect(confirmButton).toBeEnabled());
+ });
+
it.each([
{
entry: "the task instance page",
diff --git
a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx
b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx
index ddc6d22acd3..b97a9a6247e 100644
---
a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx
+++
b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx
@@ -201,7 +201,6 @@ const ClearTaskInstanceDialog = (props: Props) => {
closeDialog();
};
- // Get current DAG's bundle version to compare with task instance's DAG
version bundle version
const { data: dagDetails } = useDagServiceGetDagDetails({ dagId },
undefined, { enabled: singleRun });
const { data: dagRun } = useDagRunServiceGetDagRun({ dagId, dagRunId },
undefined, {
@@ -227,7 +226,7 @@ const ClearTaskInstanceDialog = (props: Props) => {
getRunOnLatestVersionState({
latestBundleVersion: dagDetails?.bundle_version,
latestDagVersionNumber: dagDetails?.latest_dag_version?.version_number,
- selectedBundleVersion: selectedVersion?.bundle_version,
+ selectedBundleVersion: dagRun?.bundle_version,
selectedDagVersionNumber: selectedVersion?.version_number,
selectedVersionMissing: dagRun?.dag_versions.length === 0,
});
@@ -375,7 +374,11 @@ const ClearTaskInstanceDialog = (props: Props) => {
<Modal
footerActions={
<Button
- disabled={data.total_entries === 0 || keptTaskInstances.length ===
0}
+ disabled={
+ (singleRun && (!dagRun || !dagDetails)) ||
+ data.total_entries === 0 ||
+ keptTaskInstances.length === 0
+ }
loading={isPending}
onClick={gateRequest === undefined ? confirmClear : onOpen}
>
diff --git a/airflow-core/tests/unit/executors/test_workloads.py
b/airflow-core/tests/unit/executors/test_workloads.py
index dafcbbdede6..68bdb914744 100644
--- a/airflow-core/tests/unit/executors/test_workloads.py
+++ b/airflow-core/tests/unit/executors/test_workloads.py
@@ -458,6 +458,7 @@ class TestExecuteTaskMakeVersionData:
ti.dag_version.version_data = ti_dag_version_data
if has_created_dag_version:
+ ti.dag_run.created_dag_version.bundle_version = bundle_version
ti.dag_run.created_dag_version.version_data = version_data
else:
ti.dag_run.created_dag_version = None
@@ -533,6 +534,7 @@ class TestExecuteCallbackMakeVersionData:
dag_run.dag_model.bundle_name = "test-bundle"
dag_run.dag_model.relative_fileloc = "dags/test_dag.py"
if has_created_dag_version:
+ dag_run.created_dag_version.bundle_version = bundle_version
dag_run.created_dag_version.version_data = version_data
else:
dag_run.created_dag_version = None
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 9363b694775..a90b8bd00fa 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -10496,10 +10496,12 @@ class TestSchedulerJob:
[
pytest.param(None, "abc123-sha", None,
id="disable_bundle_versioning"),
pytest.param("abc123-sha", "abc123-sha", "abc123-sha",
id="versioning_enabled"),
+ pytest.param("old-sha", "abc123-sha", "old-sha",
id="bundle_only_update"),
],
)
+ @pytest.mark.parametrize("request_type", [TaskCallbackRequest,
EmailRequest])
def test_external_kill_callback_bundle_version_follows_dag_run(
- self, dag_maker, session, dag_run_bv, dag_version_bv, expected_bv
+ self, dag_maker, session, dag_run_bv, dag_version_bv, expected_bv,
request_type
):
"""
TaskCallbackRequest.bundle_version must mirror dag_run.bundle_version,
not
@@ -10509,7 +10511,12 @@ class TestSchedulerJob:
on-disk code as the task did.
"""
with dag_maker(dag_id=f"ext_kill_bv_{dag_run_bv or 'none'}",
fileloc="/test_path1/"):
- EmptyOperator(task_id="t1", on_failure_callback=lambda ctx: None)
+ EmptyOperator(
+ task_id="t1",
+ on_failure_callback=(lambda ctx: None) if request_type is
TaskCallbackRequest else None,
+ email="[email protected]" if request_type is EmailRequest else
None,
+ email_on_failure=True,
+ )
dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
ti = dr.get_task_instance(task_id="t1", session=session)
@@ -10531,27 +10538,80 @@ class TestSchedulerJob:
self.job_runner.executor.callback_sink.send.assert_called_once()
request = self.job_runner.executor.callback_sink.send.call_args[0][0]
- assert isinstance(request, TaskCallbackRequest)
+ assert isinstance(request, request_type)
assert request.bundle_version == expected_bv
- def test_heartbeat_timeout_callback_bundle_version_follows_dag_run(self,
dag_maker, session):
+ @pytest.mark.parametrize("request_type", [TaskCallbackRequest,
EmailRequest])
+ def test_external_kill_callback_version_data_follows_dag_run(self,
dag_maker, session, request_type):
+ """
+ The callback must ship the manifest of the run's pinned Dag version,
matching
+ bundle_version, even when the TI was moved to a newer Dag version
mid-run.
+ """
+ with dag_maker(dag_id="ext_kill_version_data", fileloc="/test_path1/"):
+ EmptyOperator(
+ task_id="t1",
+ on_failure_callback=(lambda ctx: None) if request_type is
TaskCallbackRequest else None,
+ email="[email protected]" if request_type is EmailRequest else
None,
+ email_on_failure=True,
+ )
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+
+ run_manifest = {"files": {"dags/my_dag.py": "v1-object-id"}}
+ run_dag_version = dr.created_dag_version
+ run_dag_version.bundle_version = "v1-sha"
+ run_dag_version.version_data = run_manifest
+ newer_dag_version = DagVersion(
+ dag_id=dr.dag_id,
+ version_number=run_dag_version.version_number + 1,
+ bundle_name=run_dag_version.bundle_name,
+ bundle_version="v2-sha",
+ version_data={"files": {"dags/my_dag.py": "v2-object-id"}},
+ )
+ session.add(newer_dag_version)
+ session.flush()
+ dr.bundle_version = "v1-sha"
+ ti = dr.get_task_instance(task_id="t1", session=session)
+ ti.dag_version_id = newer_dag_version.id
+ ti.state = State.QUEUED
+ session.merge(run_dag_version)
+ session.merge(dr)
+ session.merge(ti)
+ session.commit()
+
+ executor = MockExecutor(do_update=False)
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
+ executor.event_buffer[TaskInstanceUuid(ti.id)] = State.FAILED, None
+
+ self.job_runner._process_executor_events(executor=executor,
session=session)
+
+ self.job_runner.executor.callback_sink.send.assert_called_once()
+ request = self.job_runner.executor.callback_sink.send.call_args[0][0]
+ assert isinstance(request, request_type)
+ assert request.bundle_version == "v1-sha"
+ assert request.version_data == run_manifest
+
+ @pytest.mark.parametrize("bundle_version", [None, "old-sha"])
+ @time_machine.travel(DEFAULT_DATE, tick=False)
+ def test_heartbeat_timeout_callback_bundle_version_follows_dag_run(
+ self, dag_maker, session, bundle_version
+ ):
"""
Same invariant as the external-kill path, exercised through
_find_and_purge_task_instances_without_heartbeats.
"""
with dag_maker(dag_id="hb_timeout_bv", fileloc="/test_path1/"):
- EmptyOperator(task_id="t1")
+ EmptyOperator(task_id="t1", email="[email protected]",
email_on_failure=True)
dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
- executor = MagicMock()
+ executor = MockExecutor(do_update=False)
scheduler_job = Job()
self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
ti = dr.get_task_instance(task_id="t1", session=session)
dag_version = ti.dag_version
- # disable_bundle_versioning state: DagVersion has a SHA, dag_run is
unpinned.
dag_version.bundle_version = "abc123-sha"
- dr.bundle_version = None
+ dr.bundle_version = bundle_version
ti.state = TaskInstanceState.RUNNING
ti.queued_by_job_id = scheduler_job.id
ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
@@ -10562,10 +10622,10 @@ class TestSchedulerJob:
self.job_runner._find_and_purge_task_instances_without_heartbeats()
- executor.send_callback.assert_called_once()
- request = executor.send_callback.call_args[0][0]
- assert isinstance(request, TaskCallbackRequest)
- assert request.bundle_version is None
+ requests = [call.args[0] for call in
executor.callback_sink.send.call_args_list]
+ assert len(requests) == 2
+ assert {type(request) for request in requests} ==
{TaskCallbackRequest, EmailRequest}
+ assert all(request.bundle_version == bundle_version for request in
requests)
@time_machine.travel(DEFAULT_DATE, tick=False)
def test_heartbeat_timeout_preserves_failure_email(self, dag_maker,
session):
@@ -10886,8 +10946,10 @@ class TestSchedulerJob:
assert ti_lock_calls[0].kwargs["session"] is session
@conf_vars({("scheduler", "num_stuck_in_queued_retries"): "1"})
+ @pytest.mark.parametrize("bundle_version", [None, "old-sha"])
+ @time_machine.travel(DEFAULT_DATE, tick=False)
def test_stuck_in_queued_callback_bundle_version_follows_dag_run(
- self, dag_maker, session, mock_executors
+ self, dag_maker, session, mock_executors, bundle_version
):
"""
Same invariant as the external-kill path, exercised through
@@ -10902,7 +10964,7 @@ class TestSchedulerJob:
ti = dr.get_task_instance(task_id="op1", session=session)
dag_version = ti.dag_version
dag_version.bundle_version = "abc123-sha"
- dr.bundle_version = None
+ dr.bundle_version = bundle_version
ti.state = State.QUEUED
ti.queued_dttm = timezone.utcnow()
session.merge(dag_version)
@@ -10928,7 +10990,7 @@ class TestSchedulerJob:
mock_executors[0].send_callback.assert_called_once()
request = mock_executors[0].send_callback.call_args[0][0]
assert isinstance(request, TaskCallbackRequest)
- assert request.bundle_version is None
+ assert request.bundle_version == bundle_version
def test_stuck_in_queued_failure_callback_for_removed_mapped_task(
self, dag_maker, session, mock_executors
diff --git a/airflow-core/tests/unit/models/test_dag_version.py
b/airflow-core/tests/unit/models/test_dag_version.py
index 69c1703da2d..e5042be9dc8 100644
--- a/airflow-core/tests/unit/models/test_dag_version.py
+++ b/airflow-core/tests/unit/models/test_dag_version.py
@@ -274,11 +274,17 @@ class TestResolveVersionData:
("dag_version", "bundle_version", "expected"),
[
pytest.param(
- mock.Mock(version_data={"schema_version": 1}),
+ mock.Mock(bundle_version="abc123",
version_data={"schema_version": 1}),
"abc123",
{"schema_version": 1},
id="pinned-with-data",
),
+ pytest.param(
+ mock.Mock(bundle_version="def456",
version_data={"schema_version": 1}),
+ "abc123",
+ None,
+ id="pinned-to-overwritten-bundle-version",
+ ),
pytest.param(
mock.Mock(version_data={"schema_version": 1}),
None,
diff --git a/airflow-core/tests/unit/models/test_trigger.py
b/airflow-core/tests/unit/models/test_trigger.py
index 54a89a8ce69..d03ccd58330 100644
--- a/airflow-core/tests/unit/models/test_trigger.py
+++ b/airflow-core/tests/unit/models/test_trigger.py
@@ -36,6 +36,7 @@ from airflow.jobs.triggerer_job_runner import
TriggererJobRunner
from airflow.models import TaskInstance, Trigger
from airflow.models.asset import AssetEvent, AssetModel, AssetWatcherModel
from airflow.models.callback import Callback, TriggererCallback
+from airflow.models.dag_version import DagVersion
from airflow.models.trigger import handle_event_submit
from airflow.models.xcom import XComModel
from airflow.providers.standard.operators.empty import EmptyOperator
@@ -690,8 +691,7 @@ def test_submit_event_task_end(mock_utcnow, session,
create_task_instance, event
@patch("airflow.callbacks.database_callback_sink.DatabaseCallbackSink.send")
def test_submit_event_task_end_callback_includes_version_data(mock_send,
session, create_task_instance):
- """A finished deferred task's callback carries version_data for a pinned
run, and its
- bundle_version is derived from the same dag_version so the two cannot
diverge."""
+ """A finished deferred task's callback carries version_data for a pinned
run."""
version_data = {"schema_version": 1, "files": {"dags/my_dag.py": "ver123"}}
trigger = Trigger(classpath="does.not.matter", kwargs={})
@@ -715,6 +715,65 @@ def
test_submit_event_task_end_callback_includes_version_data(mock_send, session
assert request.version_data == version_data
+@patch("airflow.callbacks.database_callback_sink.DatabaseCallbackSink.send")
+def test_submit_event_task_end_callback_uses_run_bundle_after_code_only_update(
+ mock_send, session, create_task_instance
+):
+ """A code-only update advances the reused Dag version's bundle while the
run stays pinned."""
+ trigger = Trigger(classpath="does.not.matter", kwargs={})
+ session.add(trigger)
+ task_instance = create_task_instance(
+ session=session, logical_date=timezone.utcnow(), state=State.DEFERRED
+ )
+ task_instance.trigger_id = trigger.id
+ task_instance.dag_run.bundle_version = "v1-sha"
+ task_instance.dag_version.bundle_version = "v2-sha"
+ session.commit()
+
+ Trigger.submit_event(trigger.id, TaskSuccessEvent(), session=session)
+ session.flush()
+
+ mock_send.assert_called_once()
+ assert mock_send.call_args.kwargs["callback"].bundle_version == "v1-sha"
+
+
+@patch("airflow.callbacks.database_callback_sink.DatabaseCallbackSink.send")
+def
test_submit_event_task_end_callback_uses_run_manifest_when_ti_on_newer_version(
+ mock_send, session, create_task_instance
+):
+ """The callback ships the run's pinned bundle and manifest after the TI
moved to a newer Dag version."""
+ trigger = Trigger(classpath="does.not.matter", kwargs={})
+ session.add(trigger)
+ task_instance = create_task_instance(
+ session=session, logical_date=timezone.utcnow(), state=State.DEFERRED
+ )
+ task_instance.trigger_id = trigger.id
+ run_manifest = {"files": {"dags/my_dag.py": "v1-object-id"}}
+ run_dag_version = task_instance.dag_run.created_dag_version
+ run_dag_version.bundle_version = "v1-sha"
+ run_dag_version.version_data = run_manifest
+ newer_dag_version = DagVersion(
+ dag_id=task_instance.dag_id,
+ version_number=run_dag_version.version_number + 1,
+ bundle_name=run_dag_version.bundle_name,
+ bundle_version="v2-sha",
+ version_data={"files": {"dags/my_dag.py": "v2-object-id"}},
+ )
+ session.add(newer_dag_version)
+ session.flush()
+ task_instance.dag_run.bundle_version = "v1-sha"
+ task_instance.dag_version_id = newer_dag_version.id
+ session.commit()
+
+ Trigger.submit_event(trigger.id, TaskSuccessEvent(), session=session)
+ session.flush()
+
+ mock_send.assert_called_once()
+ request = mock_send.call_args.kwargs["callback"]
+ assert request.bundle_version == "v1-sha"
+ assert request.version_data == run_manifest
+
+
@pytest.mark.parametrize(
("retries", "expected_state", "expected_callback_type",
"expect_history_row"),
[
