diff --git a/airflow-core/newsfragments/69842.bugfix.rst b/airflow-core/newsfragments/69842.bugfix.rst new file mode 100644 index 0000000000000..8d61335118517 --- /dev/null +++ b/airflow-core/newsfragments/69842.bugfix.rst @@ -0,0 +1 @@ +Keep the triggerer running when a deferred task is missing from its serialized Dag. diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 355763aed5f13..27318adab63ad 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -50,6 +50,7 @@ from airflow._shared.observability.metrics import stats from airflow._shared.timezones import timezone from airflow.configuration import conf +from airflow.exceptions import TaskNotFound from airflow.executors import workloads from airflow.executors.workloads.task import TaskInstanceDTO from airflow.jobs.base_job_runner import BaseJobRunner @@ -879,12 +880,26 @@ def _create_workload( ) if serialized_dag_model: - task = serialized_dag_model.dag.get_task(trigger.task_instance.task_id) + try: + task = serialized_dag_model.dag.get_task(trigger.task_instance.task_id) + except TaskNotFound: + # Avoid this lookup for ordinary deferrals once trigger context requirements are persisted; + # tracked at https://github.com/apache/airflow/issues/69841 + log.warning( + "Task for deferred trigger was not found in serialized Dag; " + "starting trigger without Dag context", + trigger_id=trigger.id, + ti_id=trigger.task_instance.id, + dag_id=trigger.task_instance.dag_id, + task_id=trigger.task_instance.task_id, + dag_version_id=trigger.task_instance.dag_version_id, + ) + task = None # When a TaskInstance of a Trigger contains a task with start_from_trigger enabled, # it means we need to load the SerializedDagModel so we can build a RuntimeTaskInstance later on which # will allow us to build a context on which we will render the templated fields. - if task.start_from_trigger: + if task is not None and task.start_from_trigger: log.info("Start from trigger enabled for task %s", task.task_id) dag_run = trigger.task_instance.get_dagrun(session=session) diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index 0f307fc150520..3866c563e42a0 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -30,6 +30,7 @@ import uuid from collections.abc import AsyncIterator from socket import socket, socketpair +from types import SimpleNamespace from typing import TYPE_CHECKING, Any from unittest import mock from unittest.mock import ANY, AsyncMock, MagicMock, patch @@ -49,6 +50,7 @@ from structlog.typing import FilteringBoundLogger from airflow._shared.timezones import timezone +from airflow.exceptions import TaskNotFound from airflow.executors import workloads from airflow.executors.workloads.task import TaskInstanceDTO from airflow.executors.workloads.trigger import RunTrigger @@ -68,6 +70,7 @@ ) from airflow.models import Connection, DagModel, DagRun, Trigger, Variable from airflow.models.dag_version import DagVersion +from airflow.models.dagbag import DBDagBag from airflow.models.dagbundle import DagBundleModel from airflow.models.serialized_dag import SerializedDagModel from airflow.models.xcom import XComModel @@ -609,6 +612,147 @@ def test_create_workload_uses_supervisor_id_without_job(jobless_supervisor, mock assert factory.log_path == f"/logs/ti.trigger.{jobless_supervisor.id}.log" +def test_create_workload_missing_task_returns_context_free_workload(jobless_supervisor, mocker, caplog): + ti = SimpleNamespace( + id=uuid.uuid4(), + dag_id="example_dag", + dag_version_id=uuid.uuid4(), + task_id="removed_task", + trigger_timeout=None, + ) + trigger = SimpleNamespace( + id=7, + classpath="some.path.Trigger", + encrypted_kwargs="encrypted", + task_instance=ti, + ) + + ser_ti = mocker.Mock(spec=TaskInstanceDTO) + mocker.patch( + "airflow.jobs.triggerer_job_runner.TaskInstanceDTO.model_validate", + autospec=True, + return_value=ser_ti, + ) + + def get_task(task_id): + raise TaskNotFound(task_id) + + serialized_dag_model = SimpleNamespace(dag=SimpleNamespace(get_task=get_task)) + dag_bag = mocker.Mock(spec=DBDagBag) + dag_bag.get_serialized_dag_model.return_value = serialized_dag_model + + workload = jobless_supervisor._create_workload( + trigger=trigger, + dag_bag=dag_bag, + render_log_fname=lambda *, ti: "/logs/ti", + session=mocker.sentinel.session, + ) + + assert workload == RunTrigger( + id=trigger.id, + classpath=trigger.classpath, + encrypted_kwargs=trigger.encrypted_kwargs, + ti=ser_ti, + ) + assert { + "event": "Task for deferred trigger was not found in serialized Dag; " + "starting trigger without Dag context", + "trigger_id": trigger.id, + "ti_id": trigger.task_instance.id, + "dag_id": trigger.task_instance.dag_id, + "task_id": trigger.task_instance.task_id, + "dag_version_id": trigger.task_instance.dag_version_id, + } in caplog + + +def test_build_trigger_workloads_continues_after_missing_task(jobless_supervisor, mocker): + def get_missing_task(task_id): + raise TaskNotFound(task_id) + + stale_ti = SimpleNamespace( + id=uuid.uuid4(), + dag_id="example_dag", + dag_version_id=uuid.uuid4(), + task_id="removed_task", + trigger_timeout=None, + ) + stale_trigger = SimpleNamespace( + id=1, + classpath="some.path.StaleTrigger", + encrypted_kwargs="stale", + task_instance=stale_ti, + ) + healthy_ti = SimpleNamespace( + id=uuid.uuid4(), + dag_id="example_dag", + dag_version_id=uuid.uuid4(), + task_id="healthy_task", + trigger_timeout=None, + get_dagrun=lambda *, session: SimpleNamespace( + dag_run_data=SimpleNamespace( + model_dump=lambda **kwargs: {"dag_id": "example_dag"}, + ) + ), + ) + healthy_trigger = SimpleNamespace( + id=2, + classpath="some.path.HealthyTrigger", + encrypted_kwargs="healthy", + task_instance=healthy_ti, + ) + + serialized_stale_dag = SimpleNamespace(dag=SimpleNamespace(get_task=get_missing_task)) + serialized_healthy_dag = SimpleNamespace( + data={"dag": "serialized"}, + dag=SimpleNamespace( + get_task=lambda task_id: SimpleNamespace(task_id=task_id, start_from_trigger=True) + ), + ) + + dag_bag = mocker.Mock(spec=DBDagBag) + serialized_dags = { + stale_trigger.task_instance.dag_version_id: serialized_stale_dag, + healthy_trigger.task_instance.dag_version_id: serialized_healthy_dag, + } + dag_bag.get_serialized_dag_model.side_effect = lambda *, version_id, session: serialized_dags[version_id] + mocker.patch("airflow.jobs.triggerer_job_runner.DBDagBag", autospec=True, return_value=dag_bag) + mocker.patch( + "airflow.jobs.triggerer_job_runner.log_filename_template_renderer", + autospec=True, + return_value=lambda *, ti: "/logs/ti", + ) + mocker.patch( + "airflow.jobs.triggerer_job_runner.create_session", + autospec=True, + return_value=contextlib.nullcontext(mocker.sentinel.session), + ) + mocker.patch( + "airflow.jobs.triggerer_job_runner.TaskInstanceDTO.model_validate", + autospec=True, + side_effect=[mocker.Mock(spec=TaskInstanceDTO), mocker.Mock(spec=TaskInstanceDTO)], + ) + mocker.patch.object( + TriggerRunnerSupervisor, + "fetch_trigger_details", + autospec=True, + return_value={stale_trigger.id: stale_trigger, healthy_trigger.id: healthy_trigger}, + ) + mocker.patch.object( + TriggerRunnerSupervisor, + "fetch_non_task_trigger_ids", + autospec=True, + return_value=set(), + ) + + workloads = jobless_supervisor.build_trigger_workloads({stale_trigger.id, healthy_trigger.id}) + + workload_by_id = {workload.id: workload for workload in workloads} + assert workload_by_id.keys() == {stale_trigger.id, healthy_trigger.id} + assert workload_by_id[stale_trigger.id].dag_data is None + assert workload_by_id[healthy_trigger.id].dag_data == serialized_healthy_dag.data + assert workload_by_id[healthy_trigger.id].dag_run_data == {"dag_id": "example_dag"} + + def test_create_workload_sets_watched_assets_for_asset_only_trigger(jobless_supervisor, mocker): """_create_workload() should populate watched_assets when trigger.task_instance is None and assets exist.""" asset1 = mocker.Mock(spec=Asset)