nprasad2077/airflow
1
0 Fork 0
mirror of https://github.com/apache/airflow.git synced 2026-11-04 07:12:48 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
63b615b252589ef49b8504f2dbe2942b543ec896
airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py
T
Niko Oliveira 63b615b252 Scope plugin macros to their team in multi-team mode (#73224) 
A team-scoped plugin's macros were reachable from every task's templates,
so any team could call macros written to talk to another team's systems.

The worker decides this from a new ``multi_team`` flag on the run context
rather than from its own config, which is not guaranteed to match the
scheduler's (misconfigured or intentional attempt to read macros from
other teams): reading it as disabled while it is in fact enabled would drop
scoping and apply a team's macros everywhere. The existing ``team_name``
cannot carry that signal, being ``None`` both for a teamless task and for
every task when multi-team is off, so a plugin declaring a team would
otherwise lose its macros in single-team deployments.

Scoping is applied at attribute access rather than when macros are loaded,
because the loader mutates a process-wide module and a worker process can
serve tasks from more than one team.
2026-09-22 09:43:16 -07:00

2526 lines
101 KiB
Python
Raw Blame History

#
# 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.
"""The entrypoint for the actual task execution process."""
from __future__ import annotations
import contextvars
import functools
import inspect
import os
import sys
import time
from collections.abc import Callable, Iterable, Iterator, Mapping
from contextlib import ExitStack, contextmanager, suppress
from datetime import datetime, timedelta, timezone
from itertools import product
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Any, Literal
from urllib.parse import quote
import attrs
import lazy_object_proxy
import structlog
from opentelemetry import trace
from opentelemetry.trace import INVALID_SPAN, Status, StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from pydantic import AwareDatetime, ConfigDict, Field, JsonValue, TypeAdapter
from structlog.contextvars import bind_contextvars
from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersionLock
from airflow.dag_processing.bundles.manager import DagBundlesManager
from airflow.sdk._shared.observability.metrics import stats
from airflow.sdk._shared.observability.metrics.stats import build_dag_metric_tags
from airflow.sdk._shared.observability.traces import get_task_span_detail_level
from airflow.sdk._shared.template_rendering import truncate_rendered_value
from airflow.sdk.api.client import get_hostname, getuser
from airflow.sdk.api.datamodels._generated import (
AssetProfile,
DagRun,
PreviousTIResponse,
TaskInstance,
TaskInstanceState,
TIRunContext,
)
from airflow.sdk.bases.operator import BaseOperator, ExecutorSafeguard
from airflow.sdk.bases.skipmixin import XCOM_SKIPMIXIN_KEY
from airflow.sdk.bases.xcom import BaseXCom
from airflow.sdk.configuration import conf
from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager
from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set
from airflow.sdk.definitions.asset import (
Asset,
AssetAlias,
AssetNameRef,
AssetUniqueKey,
AssetUriRef,
)
from airflow.sdk.definitions.mappedoperator import MappedOperator
from airflow.sdk.definitions.param import process_params
from airflow.sdk.exceptions import (
AirflowException,
AirflowFailException,
AirflowInactiveAssetInInletOrOutletException,
AirflowRescheduleException,
AirflowRuntimeError,
AirflowSensorTimeout,
AirflowTaskTerminated,
AirflowTaskTimeout,
ErrorType,
TaskAwaitingInput,
TaskDeferred,
)
from airflow.sdk.execution_time.callback_runner import create_executable_runner
from airflow.sdk.execution_time.comms import (
AssetEventDagRunReferenceResult,
AwaitInputTask,
CommsDecoder,
DagResult,
DagRunStateResult,
DeferTask,
DRCount,
ErrorResponse,
GetDag,
GetDagRunState,
GetDRCount,
GetPreviousDagRun,
GetPreviousTI,
GetTaskBreadcrumbs,
GetTaskRescheduleStartDate,
GetTaskStates,
GetTICount,
InactiveAssetsResult,
PreviousDagRunResult,
PreviousTIResult,
RescheduleTask,
ResendLoggingFD,
RetryTask,
SentFDs,
SetRenderedFields,
SetRenderedMapIndex,
SkipDownstreamTasks,
StartupDetails,
SucceedTask,
TaskBreadcrumbsResult,
TaskRescheduleStartDate,
TaskState,
TaskStatesResult,
TICount,
ToSupervisor,
ToTask,
TriggerDagRun,
ValidateInletsAndOutlets,
)
from airflow.sdk.execution_time.context import (
AssetStateStoreAccessors,
ConnectionAccessor,
InletEventsAccessors,
MacrosAccessor,
OutletEventAccessors,
TaskStateStoreAccessor,
TriggeringAssetEventsAccessor,
VariableAccessor,
_get_worker_state_store_backend,
context_get_outlet_events,
context_to_airflow_vars,
get_previous_dagrun_success,
set_current_context,
)
from airflow.sdk.execution_time.email_backend import (
_DEFAULT_EMAIL_BACKEND,
_ErrorEmailNotifier,
_LegacyEmailBackendNotifier,
)
from airflow.sdk.execution_time.sentry import Sentry
from airflow.sdk.execution_time.xcom import XCom
from airflow.sdk.listener import get_listener_manager
from airflow.sdk.observability.metrics import stats_utils
from airflow.sdk.serde import allow_class, iter_pydantic_models
from airflow.sdk.state import TaskScope
from airflow.sdk.timezone import coerce_datetime
if TYPE_CHECKING:
import jinja2
from pendulum.datetime import DateTime
from structlog.typing import FilteringBoundLogger as Logger
from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator
from airflow.sdk.definitions.context import Context
from airflow.sdk.definitions.retry_policy import RetryDecision
from airflow.sdk.exceptions import DagRunTriggerException
from airflow.sdk.types import OutletEventAccessorsProtocol
log = structlog.get_logger("task")
tracer = trace.get_tracer(__name__)
class detail_span:
"""Context manager and decorator that creates a child span when detail level > 1."""
def __init__(self, *args, **kwargs):
self._args = args
self._kwargs = kwargs
self._ctx = None
def _make_ctx(self):
parent_span = trace.get_current_span()
config_level = get_task_span_detail_level(span=parent_span)
if config_level > 1:
return tracer.start_as_current_span(*self._args, **self._kwargs)
return trace.INVALID_SPAN
def __enter__(self):
self._ctx = self._make_ctx()
return self._ctx.__enter__()
def __exit__(self, *exc_info):
return self._ctx.__exit__(*exc_info)
def __call__(self, f):
@functools.wraps(f)
def wrapper(*inner_args, **inner_kwargs):
with self._make_ctx():
return f(*inner_args, **inner_kwargs)
wrapper.__signature__ = inspect.signature(f)
return wrapper
@contextmanager
def _make_task_span(msg: StartupDetails):
parent_context = (
TraceContextTextMapPropagator().extract(msg.ti.context_carrier) if msg.ti.context_carrier else None
)
ti = msg.ti
span_name = f"worker.{ti.task_id}"
if ti.map_index is not None and ti.map_index >= 0:
span_name += f"[{ti.map_index}]"
with tracer.start_as_current_span(span_name, context=parent_context) as span:
span.set_attributes(
{
"airflow.dag_id": ti.dag_id,
"airflow.task_id": ti.task_id,
"airflow.dag_run.run_id": ti.run_id,
"airflow.task_instance.try_number": ti.try_number,
"airflow.task_instance.map_index": ti.map_index if ti.map_index is not None else -1,
}
)
yield span
class TaskRunnerMarker:
"""Marker for listener hooks, to properly detect from which component they are called."""
# TODO: Move this entire class into a separate file:
# `airflow/sdk/execution_time/task_instance.py`
# or `airflow/sdk/execution_time/runtime_ti.py`
class RuntimeTaskInstance(TaskInstance):
model_config = ConfigDict(arbitrary_types_allowed=True)
task: BaseOperator
bundle_instance: BaseDagBundle
_cached_template_context: Context | None = None
"""The Task Instance context. This is used to cache get_template_context."""
_terminal_state_send_failed: bool = False
"""True when the supervisor IPC send for a non-success terminal state raised; signals main() to sys.exit(1) after finalize() so the supervisor doesn't misclassify the run as SUCCESS via exit code 0."""
_failure_metrics_emitted: bool = False
"""True once the failure counters have been recorded, so a second pass through the failure path (one attempt raised part-way through) does not count the same failure twice."""
_ti_context_from_server: Annotated[TIRunContext | None, Field(repr=False)] = None
"""The Task Instance context from the API server, if any."""
max_tries: int = 0
"""The maximum number of retries for the task."""
start_date: AwareDatetime
"""Start date of the task instance."""
end_date: AwareDatetime | None = None
state: TaskInstanceState | None = None
is_mapped: bool | None = None
"""True if the original task was mapped."""
rendered_map_index: str | None = None
sentry_integration: str = ""
@property
def stats_tags(self) -> dict[str, str]:
"""Metric tags for this task instance, including dag tags and team_name when available."""
tags: dict[str, str] = {}
if conf.getboolean("metrics", "dag_tags_in_metrics", fallback=False):
tags.update(build_dag_metric_tags(self.task.dag.tags))
# Built-in keys always win on collision.
tags.update(dag_id=self.dag_id, task_id=self.task_id)
if self._ti_context_from_server:
# run_type keeps the tag set consistent with the scheduler-side TaskInstance.stats_tags.
# Coerce the DagRunType enum to its bare value so it serializes as e.g. "scheduled" rather
# than "dagruntype.scheduled" (matching the scheduler, which emits the plain string).
run_type = self._ti_context_from_server.dag_run.run_type
tags["run_type"] = getattr(run_type, "value", run_type)
if self._ti_context_from_server.dag_run.team_name:
tags["team_name"] = self._ti_context_from_server.dag_run.team_name
return tags
def __rich_repr__(self):
yield "id", self.id
yield "task_id", self.task_id
yield "dag_id", self.dag_id
yield "run_id", self.run_id
yield "max_tries", self.max_tries
yield "task", type(self.task)
yield "start_date", self.start_date
__rich_repr__.angular = True # type: ignore[attr-defined]
@detail_span("get_template_context")
def get_template_context(self) -> Context:
# TODO: Move this to `airflow.sdk.execution_time.context`
# once we port the entire context logic from airflow/utils/context.py ?
from airflow.sdk.plugins_manager import integrate_macros_plugins
integrate_macros_plugins()
dag_run_conf: dict[str, Any] | None = None
macros_accessor = MacrosAccessor()
if from_server := self._ti_context_from_server:
dag_run_conf = from_server.dag_run.conf or dag_run_conf
macros_accessor = MacrosAccessor(
team_name=from_server.dag_run.team_name, multi_team=bool(from_server.multi_team)
)
validated_params = process_params(self.task.dag, self.task, dag_run_conf, suppress_exception=False)
# Cache the context object, which ensures that all calls to get_template_context
# are operating on the same context object.
if self._cached_template_context is None:
self._cached_template_context = {
# From the Task Execution interface
"dag": self.task.dag,
"inlets": self.task.inlets,
"map_index_template": self.task.map_index_template,
"outlets": self.task.outlets,
"run_id": self.run_id,
"task": self.task,
"task_instance": self,
"ti": self,
"outlet_events": OutletEventAccessors(),
"inlet_events": InletEventsAccessors(self.task.inlets),
"macros": macros_accessor,
"params": validated_params,
# TODO: Make this go through Public API longer term.
# "test_mode": task_instance.test_mode,
"var": {
"json": VariableAccessor(deserialize_json=True),
"value": VariableAccessor(deserialize_json=False),
},
"conn": ConnectionAccessor(),
"task_state_store": TaskStateStoreAccessor(
ti_id=self.id,
scope=TaskScope(
dag_id=self.dag_id,
run_id=self.run_id,
task_id=self.task_id,
map_index=self.map_index if self.map_index is not None else -1,
),
),
}
_asset_types = (Asset, AssetNameRef, AssetUriRef, AssetAlias)
if any(isinstance(i, _asset_types) for i in self.task.inlets + self.task.outlets):
self._cached_template_context["asset_state_store"] = AssetStateStoreAccessors(
self.task.inlets, self.task.outlets
)
# AssetAlias inlets are resolved to their concrete assets at context build time via
# GetAssetsByAlias comms. If an alias maps to no active assets, it doesn't contribute to
# asset_state. AssetAlias outlets are skipped downstream in context.py
if TYPE_CHECKING:
assert self._cached_template_context is not None
if from_server:
dag_run = from_server.dag_run
context_from_server: Context = {
# TODO: Assess if we need to pass these through timezone.coerce_datetime
"dag_run": dag_run, # type: ignore[typeddict-item] # Removable after #46522
"partition_key": dag_run.partition_key,
"partition_date": coerce_datetime(dag_run.partition_date),
"triggering_asset_events": TriggeringAssetEventsAccessor.build(
AssetEventDagRunReferenceResult.from_asset_event_dag_run_reference(event)
for event in dag_run.consumed_asset_events
),
"task_instance_key_str": f"{self.task.dag_id}__{self.task.task_id}__{dag_run.run_id}",
"task_reschedule_count": from_server.task_reschedule_count or 0,
"prev_start_date_success": lazy_object_proxy.Proxy(
lambda: coerce_datetime(get_previous_dagrun_success(self.id).start_date)
),
"prev_end_date_success": lazy_object_proxy.Proxy(
lambda: coerce_datetime(get_previous_dagrun_success(self.id).end_date)
),
}
self._cached_template_context.update(context_from_server)
if logical_date := coerce_datetime(dag_run.logical_date):
if TYPE_CHECKING:
assert isinstance(logical_date, DateTime)
ds = logical_date.strftime("%Y-%m-%d")
ds_nodash = ds.replace("-", "")
ts = logical_date.isoformat()
ts_nodash = logical_date.strftime("%Y%m%dT%H%M%S")
ts_nodash_with_tz = ts.replace("-", "").replace(":", "")
# logical_date and data_interval either coexist or be None together
self._cached_template_context.update(
{
# keys that depend on logical_date
"logical_date": logical_date,
"ds": ds,
"ds_nodash": ds_nodash,
"task_instance_key_str": f"{self.task.dag_id}__{self.task.task_id}__{ds_nodash}",
"ts": ts,
"ts_nodash": ts_nodash,
"ts_nodash_with_tz": ts_nodash_with_tz,
# keys that depend on data_interval
"data_interval_end": coerce_datetime(dag_run.data_interval_end),
"data_interval_start": coerce_datetime(dag_run.data_interval_start),
"prev_data_interval_start_success": lazy_object_proxy.Proxy(
lambda: coerce_datetime(get_previous_dagrun_success(self.id).data_interval_start)
),
"prev_data_interval_end_success": lazy_object_proxy.Proxy(
lambda: coerce_datetime(get_previous_dagrun_success(self.id).data_interval_end)
),
}
)
# Backward compatibility: old servers may still send upstream_map_indexes
upstream_map_indexes = getattr(from_server, "upstream_map_indexes", None)
if upstream_map_indexes is not None:
setattr(self, "_upstream_map_indexes", upstream_map_indexes)
return self._cached_template_context
@detail_span("render_templates")
def render_templates(
self, context: Context | None = None, jinja_env: jinja2.Environment | None = None
) -> BaseOperator:
"""
Render templates in the operator fields.
If the task was originally mapped, this may replace ``self.task`` with
the unmapped, fully rendered BaseOperator. The original ``self.task``
before replacement is returned.
"""
if not context:
context = self.get_template_context()
original_task = self.task
if TYPE_CHECKING:
assert context
ti = context["ti"]
if TYPE_CHECKING:
assert original_task
assert self.task
assert ti.task
# If self.task is mapped, this call replaces self.task to point to the
# unmapped BaseOperator created by this function! This is because the
# MappedOperator is useless for template rendering, and we need to be
# able to access the unmapped task instead.
self.task.render_template_fields(context, jinja_env)
self.is_mapped = original_task.is_mapped
return original_task
def _normalize_xcom_pull_params(
self,
task_ids: str | Iterable[str] | None,
dag_id: str | None,
map_indexes: int | Iterable[int] | None | ArgNotSet,
run_id: str | None,
) -> tuple[str, str, list[str], bool, bool, Iterable[int | None] | ArgNotSet]:
"""
Normalize and validate the parameters shared by ``xcom_pull`` and ``axcom_pull``.
:returns: A tuple of ``(dag_id, run_id, task_ids, single_task_requested,
single_map_index_requested, map_indexes_iterable)``. When
``map_indexes_iterable`` is ``NOTSET`` the caller should fetch *all*
map-indexes (i.e. use ``get_all`` / ``aget_all``); otherwise it is an
iterable of explicit map indexes to fetch one-by-one.
"""
if dag_id is None:
dag_id = self.dag_id
if run_id is None:
run_id = self.run_id
single_task_requested = isinstance(task_ids, (str, type(None)))
single_map_index_requested = isinstance(map_indexes, (int, type(None)))
if task_ids is None:
task_ids = [self.task_id]
elif isinstance(task_ids, str):
task_ids = [task_ids]
else:
task_ids = list(task_ids)
if not is_arg_set(map_indexes):
map_indexes_iterable: Iterable[int | None] | ArgNotSet = NOTSET
elif isinstance(map_indexes, int) or map_indexes is None:
map_indexes_iterable = [map_indexes]
elif isinstance(map_indexes, Iterable):
map_indexes_iterable = map_indexes
else:
raise TypeError(
f"Invalid type for map_indexes: expected int, iterable of ints, or None, got {type(map_indexes)}"
)
return (
dag_id,
run_id,
task_ids,
single_task_requested,
single_map_index_requested,
map_indexes_iterable,
)
def xcom_pull(
self,
task_ids: str | Iterable[str] | None = None,
dag_id: str | None = None,
key: str = BaseXCom.XCOM_RETURN_KEY,
include_prior_dates: bool = False,
*,
map_indexes: int | Iterable[int] | None | ArgNotSet = NOTSET,
default: Any = None,
run_id: str | None = None,
) -> Any:
"""
Pull XComs either from the API server (BaseXCom) or from the custom XCOM backend if configured.
The pull can be filtered optionally by certain criterion.
:param key: A key for the XCom. If provided, only XComs with matching
keys will be returned. The default key is ``'return_value'``, also
available as constant ``XCOM_RETURN_KEY``. This key is automatically
given to XComs returned by tasks (as opposed to being pushed
manually).
:param task_ids: Only XComs from tasks with matching ids will be
pulled. If *None* (default), the task_id of the calling task is used.
:param dag_id: If provided, only pulls XComs from this Dag. If *None*
(default), the Dag of the calling task is used.
:param map_indexes: If provided, only pull XComs with matching indexes.
If *None* (default), this is inferred from the task(s) being pulled
(see below for details).
:param include_prior_dates: If False, only XComs from the current
logical_date are returned. If *True*, XComs from previous dates
are returned as well.
:param run_id: If provided, only pulls XComs from a DagRun w/a matching run_id.
If *None* (default), the run_id of the calling task is used.
When pulling one single task (``task_id`` is *None* or a str) without
specifying ``map_indexes``, the return value is a single XCom entry
(map_indexes is set to map_index of the calling task instance).
When pulling task is mapped the specified ``map_index`` is used, so by default
pulling on mapped task will result in no matching XComs if the task instance
of the method call is not mapped. Otherwise, the map_index of the calling task
instance is used. Setting ``map_indexes`` to *None* will pull XCom as it would
from a non mapped task.
In either case, ``default`` (*None* if not specified) is returned if no
matching XComs are found.
When pulling multiple tasks (i.e. either ``task_id`` or ``map_index`` is
a non-str iterable), a list of matching XComs is returned. Elements in
the list is ordered by item ordering in ``task_id`` and ``map_index``.
"""
dag_id, run_id, task_ids, single_task_requested, single_map_index_requested, map_indexes_iterable = (
self._normalize_xcom_pull_params(task_ids, dag_id, map_indexes, run_id)
)
xcoms: list[Any] = []
if not is_arg_set(map_indexes_iterable):
# map_indexes was not specified — fetch all map indexes for each task
for t_id in task_ids:
values = XCom.get_all(
run_id=run_id,
key=key,
task_id=t_id,
dag_id=dag_id,
include_prior_dates=include_prior_dates,
)
xcoms.append(None) if values is None else xcoms.extend(values)
# For a single task pulling from an unmapped task, return a single value
if single_task_requested and len(xcoms) == 1:
return xcoms[0]
return xcoms
for t_id, m_idx in product(task_ids, map_indexes_iterable):
value = XCom.get_one(
run_id=run_id,
key=key,
task_id=t_id,
dag_id=dag_id,
map_index=m_idx,
include_prior_dates=include_prior_dates,
)
xcoms.append(default if value is None else value)
if single_task_requested and single_map_index_requested:
return xcoms[0]
return xcoms
async def axcom_pull(
self,
task_ids: str | Iterable[str] | None = None,
dag_id: str | None = None,
key: str = BaseXCom.XCOM_RETURN_KEY,
include_prior_dates: bool = False,
*,
map_indexes: int | Iterable[int] | None | ArgNotSet = NOTSET,
default: Any = None,
run_id: str | None = None,
) -> Any:
"""Async version of :meth:`xcom_pull`; see that method for full documentation."""
dag_id, run_id, task_ids, single_task_requested, single_map_index_requested, map_indexes_iterable = (
self._normalize_xcom_pull_params(task_ids, dag_id, map_indexes, run_id)
)
xcoms: list[Any] = []
if not is_arg_set(map_indexes_iterable):
# map_indexes was not specified — fetch all map indexes for each task
for t_id in task_ids:
values = await XCom.aget_all(
run_id=run_id,
key=key,
task_id=t_id,
dag_id=dag_id,
include_prior_dates=include_prior_dates,
)
xcoms.append(None) if values is None else xcoms.extend(values)
# For a single task pulling from an unmapped task, return a single value
if single_task_requested and len(xcoms) == 1:
return xcoms[0]
return xcoms
for t_id, m_idx in product(task_ids, map_indexes_iterable):
value = await XCom.aget_one(
run_id=run_id,
key=key,
task_id=t_id,
dag_id=dag_id,
map_index=m_idx,
include_prior_dates=include_prior_dates,
)
xcoms.append(default if value is None else value)
if single_task_requested and single_map_index_requested:
return xcoms[0]
return xcoms
def xcom_push(self, key: str, value: Any):
"""
Make an XCom available for tasks to pull.
:param key: Key to store the value under.
:param value: Value to store. Only be JSON-serializable values may be used.
"""
_xcom_push(self, key, value)
async def axcom_push(self, key: str, value: Any):
"""
Make an XCom available for tasks to pull asynchronously.
:param key: Key to store the value under.
:param value: Value to store. Only be JSON-serializable values may be used.
"""
await _axcom_push(self, key, value)
def get_relevant_upstream_map_indexes(
self, upstream: BaseOperator, ti_count: int | None, session: Any
) -> int | range | None:
"""
Compute the relevant upstream map indexes for XCom resolution.
:param upstream: The upstream operator
:param ti_count: The total count of task instances for this task's expansion
:param session: Not used (kept for API compatibility)
:return: None (use entire value), int (single index), or range (subset of indexes)
"""
from airflow.sdk.execution_time.task_mapping import get_relevant_map_indexes, get_ti_count_for_task
map_index = self.map_index
if map_index is None or map_index < 0:
return None
# If ti_count not provided, we need to query it
if ti_count is None:
ti_count = get_ti_count_for_task(self.task_id, self.dag_id, self.run_id)
if not ti_count:
return None
return get_relevant_map_indexes(
task=self.task,
run_id=self.run_id,
map_index=map_index,
ti_count=ti_count,
relative=upstream,
dag_id=self.dag_id,
)
def get_first_reschedule_date(self, context: Context) -> AwareDatetime | None:
"""Get the first reschedule date for the task instance if found, none otherwise."""
if context.get("task_reschedule_count", 0) == 0:
# If the task has not been rescheduled, there is no need to ask the supervisor
return None
max_tries: int = self.max_tries
retries: int = self.task.retries or 0
first_try_number = max_tries - retries + 1
log.debug("Requesting first reschedule date from supervisor")
response = SUPERVISOR_COMMS.send(
msg=GetTaskRescheduleStartDate(ti_id=self.id, try_number=first_try_number)
)
if TYPE_CHECKING:
assert isinstance(response, TaskRescheduleStartDate)
return response.start_date
def get_previous_dagrun(self, state: str | None = None) -> DagRun | None:
"""Return the previous Dag run before the given logical date, optionally filtered by state."""
context = self.get_template_context()
dag_run = context.get("dag_run")
log.debug("Getting previous Dag run", dag_run=dag_run)
if dag_run is None:
return None
if dag_run.logical_date is None:
return None
response = SUPERVISOR_COMMS.send(
msg=GetPreviousDagRun(dag_id=self.dag_id, logical_date=dag_run.logical_date, state=state)
)
if TYPE_CHECKING:
assert isinstance(response, PreviousDagRunResult)
return response.dag_run
def get_previous_ti(
self,
state: TaskInstanceState | None = None,
logical_date: AwareDatetime | None = None,
map_index: int = -1,
) -> PreviousTIResponse | None:
"""
Return the previous task instance matching the given criteria.
:param state: Filter by TaskInstance state
:param logical_date: Filter by logical date (returns TI before this date)
:param map_index: Filter by map_index (defaults to -1 for non-mapped tasks)
:return: Previous task instance or None if not found
"""
context = self.get_template_context()
dag_run = context.get("dag_run")
log.debug("Getting previous task instance", task_id=self.task_id, state=state)
# Use current dag run's logical_date if not provided
effective_logical_date = logical_date
if effective_logical_date is None and dag_run and dag_run.logical_date:
effective_logical_date = dag_run.logical_date
response = SUPERVISOR_COMMS.send(
msg=GetPreviousTI(
dag_id=self.dag_id,
task_id=self.task_id,
logical_date=effective_logical_date,
map_index=map_index,
state=state,
)
)
if TYPE_CHECKING:
assert isinstance(response, PreviousTIResult)
return response.task_instance
@staticmethod
def get_ti_count(
dag_id: str,
map_index: int | None = None,
task_ids: list[str] | None = None,
task_group_id: str | None = None,
logical_dates: list[datetime] | None = None,
run_ids: list[str] | None = None,
states: list[str] | None = None,
) -> int:
"""Return the number of task instances matching the given criteria."""
response = SUPERVISOR_COMMS.send(
GetTICount(
dag_id=dag_id,
map_index=map_index,
task_ids=task_ids,
task_group_id=task_group_id,
logical_dates=logical_dates,
run_ids=run_ids,
states=states,
),
)
if TYPE_CHECKING:
assert isinstance(response, TICount)
return response.count
@staticmethod
def get_task_states(
dag_id: str,
map_index: int | None = None,
task_ids: list[str] | None = None,
task_group_id: str | None = None,
logical_dates: list[datetime] | None = None,
run_ids: list[str] | None = None,
) -> dict[str, Any]:
"""Return the task states matching the given criteria."""
response = SUPERVISOR_COMMS.send(
GetTaskStates(
dag_id=dag_id,
map_index=map_index,
task_ids=task_ids,
task_group_id=task_group_id,
logical_dates=logical_dates,
run_ids=run_ids,
),
)
if TYPE_CHECKING:
assert isinstance(response, TaskStatesResult)
return response.task_states
@staticmethod
def get_task_breadcrumbs(dag_id: str, run_id: str) -> Iterable[dict[str, Any]]:
"""Return task breadcrumbs for the given dag run."""
response = SUPERVISOR_COMMS.send(GetTaskBreadcrumbs(dag_id=dag_id, run_id=run_id))
if TYPE_CHECKING:
assert isinstance(response, TaskBreadcrumbsResult)
return response.breadcrumbs
@staticmethod
def get_dr_count(
dag_id: str,
logical_dates: list[datetime] | None = None,
run_ids: list[str] | None = None,
states: list[str] | None = None,
) -> int:
"""Return the number of Dag runs matching the given criteria."""
response = SUPERVISOR_COMMS.send(
GetDRCount(
dag_id=dag_id,
logical_dates=logical_dates,
run_ids=run_ids,
states=states,
),
)
if TYPE_CHECKING:
assert isinstance(response, DRCount)
return response.count
@staticmethod
def get_dagrun_state(dag_id: str, run_id: str) -> str:
"""Return the state of the Dag run with the given Run ID."""
response = SUPERVISOR_COMMS.send(msg=GetDagRunState(dag_id=dag_id, run_id=run_id))
if TYPE_CHECKING:
assert isinstance(response, DagRunStateResult)
return response.state
@staticmethod
def get_dag(dag_id: str) -> DagResult:
"""Return the DAG with the given dag_id."""
response = SUPERVISOR_COMMS.send(msg=GetDag(dag_id=dag_id))
if TYPE_CHECKING:
assert isinstance(response, DagResult)
return response
@property
def log_url(self) -> str:
run_id = quote(self.run_id)
base_url = conf.get("api", "base_url", fallback="http://localhost:8080/")
map_index_value = self.map_index
map_index = (
f"/mapped/{map_index_value}" if map_index_value is not None and map_index_value >= 0 else ""
)
try_number_value = self.try_number
try_number = (
f"?try_number={try_number_value}" if try_number_value is not None and try_number_value > 0 else ""
)
_log_uri = f"{base_url.rstrip('/')}/dags/{self.dag_id}/runs/{run_id}/tasks/{self.task_id}{map_index}{try_number}"
return _log_uri
@property
def mark_success_url(self) -> str:
"""URL to mark TI success."""
return self.log_url
def _xcom_push(
ti: RuntimeTaskInstance,
key: str,
value: Any,
*,
mapped_length: int | None = None,
) -> None:
"""Push a XCom through XCom.set, which pushes to XCom Backend if configured."""
# Private function, as we don't want to expose the ability to manually set `mapped_length` to SDK
# consumers
if key == XCOM_SKIPMIXIN_KEY:
# The branch/skip decision is control-plane data the scheduler reads (via
# NotPreviouslySkippedDep) to skip mapped or cleared downstream tasks. It must
# bypass any custom XCom backend, which could externalize it into a pointer the
# scheduler cannot interpret, silently leaving those tasks unskipped (#50491).
_xcom_push_to_db(ti, key, value)
return
XCom.set(
key=key,
value=value,
dag_id=ti.dag_id,
task_id=ti.task_id,
run_id=ti.run_id,
map_index=ti.map_index,
dag_result=ti.task.returns_dag_result,
_mapped_length=mapped_length,
)
async def _axcom_push(
ti: RuntimeTaskInstance,
key: str,
value: Any,
*,
mapped_length: int | None = None,
) -> None:
"""Push a XCom through XCom.aset, which pushes to XCom Backend if configured."""
# Private function, as we don't want to expose the ability to manually set `mapped_length` to SDK
# consumers
await XCom.aset(
key=key,
value=value,
dag_id=ti.dag_id,
task_id=ti.task_id,
run_id=ti.run_id,
map_index=ti.map_index,
dag_result=ti.task.returns_dag_result,
_mapped_length=mapped_length,
)
def _xcom_push_to_db(ti: RuntimeTaskInstance, key: str, value: Any) -> None:
"""Push a XCom directly to metadata DB, bypassing custom xcom_backend."""
XCom._set_xcom_in_db(
key=key,
value=value,
dag_id=ti.dag_id,
task_id=ti.task_id,
run_id=ti.run_id,
map_index=ti.map_index,
)
def _maybe_reschedule_startup_failure(
*,
ti_context: TIRunContext,
log: Logger,
) -> None:
"""
Attempt to reschedule the task when a startup failure occurs.
This does not count as a retry. If the reschedule limit is exceeded, this function
returns and the caller should fail the task.
"""
missing_dag_retries = conf.getint("workers", "missing_dag_retries", fallback=3)
missing_dag_retry_delay = conf.getint("workers", "missing_dag_retry_delay", fallback=60)
reschedule_count = int(getattr(ti_context, "task_reschedule_count", 0) or 0)
if missing_dag_retries > 0 and reschedule_count < missing_dag_retries:
raise AirflowRescheduleException(
reschedule_date=datetime.now(tz=timezone.utc) + timedelta(seconds=missing_dag_retry_delay)
)
log.error(
"Startup reschedule limit exceeded",
reschedule_count=reschedule_count,
max_reschedules=missing_dag_retries,
)
def _register_deserialization_allowed_classes(dag, log: Logger) -> None:
"""
Register every operator-declared XCom model class in the deserialization allow-list.
Runs once per task-run startup, walking the whole DAG, so a consumer task can
deserialize a producer's structured output even though only the producer
constructs the value. Reads the declared classes off real operators
(``getattr``) and off not-yet-expanded mapped operators (``partial_kwargs``),
and works regardless of how the DAG was loaded -- a fresh parse, or a cache
that reconstructs operators without running ``__init__``.
Failures to register a single class are logged and skipped rather than failing
task startup; a genuinely undeserializable declaration surfaces later as the
normal allow-list ImportError at consume time.
"""
for op in dag.tasks:
if isinstance(op, MappedOperator):
fields = getattr(op.operator_class, "deserialization_allowed_class_fields", ())
values = [op.partial_kwargs.get(field) for field in fields]
else:
fields = getattr(op, "deserialization_allowed_class_fields", ())
values = [getattr(op, field, None) for field in fields]
for value in values:
if value is None:
continue
for model_cls in iter_pydantic_models(value):
try:
allow_class(model_cls)
except ValueError as exc:
log.warning(
"Skipping XCom deserialization registration for a model class",
task_id=op.task_id,
error=str(exc),
)
@detail_span("parse")
def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance:
# TODO: Task-SDK:
# Using BundleDagBag here is about 98% wrong, but it'll do for now
from airflow.dag_processing.dagbag import BundleDagBag
bundle_info = what.bundle_info
bundle_prepare_start = time.monotonic()
bundle_instance = DagBundlesManager().get_bundle(
name=bundle_info.name,
version=bundle_info.version,
version_data=bundle_info.version_data,
)
bundle_instance.initialize()
_verify_bundle_access(bundle_instance, log)
bundle_prepare_ms = int((time.monotonic() - bundle_prepare_start) * 1000)
dag_absolute_path = os.fspath(Path(bundle_instance.path, what.dag_rel_path))
dag_file_parse_start = time.monotonic()
bag = BundleDagBag(
dag_folder=dag_absolute_path,
safe_mode=False,
load_op_links=False,
bundle_path=bundle_instance.path,
bundle_name=bundle_info.name,
)
dag_file_parse_ms = int((time.monotonic() - dag_file_parse_start) * 1000)
if TYPE_CHECKING:
assert what.ti.dag_id
try:
dag = bag.dags[what.ti.dag_id]
except KeyError:
log.error(
"Dag not found during start up", dag_id=what.ti.dag_id, bundle=bundle_info, path=what.dag_rel_path
)
_maybe_reschedule_startup_failure(ti_context=what.ti_context, log=log)
sys.exit(1)
# install_loader()
try:
task = dag.task_dict[what.ti.task_id]
except KeyError:
log.error(
"Task not found in Dag during start up",
dag_id=dag.dag_id,
task_id=what.ti.task_id,
bundle=bundle_info,
path=what.dag_rel_path,
)
_maybe_reschedule_startup_failure(ti_context=what.ti_context, log=log)
sys.exit(1)
if not isinstance(task, (BaseOperator, MappedOperator)):
raise TypeError(
f"task is of the wrong type, got {type(task)}, wanted {BaseOperator} or {MappedOperator}"
)
# Register operator-declared XCom model classes (e.g. ``output_type``) for
# the whole DAG so this task can deserialize structured output produced by
# any other task -- including mapped producers and DAGs loaded from a cache
# that bypasses operator ``__init__``.
_register_deserialization_allowed_classes(dag, log)
# Surface the post-RUNNING startup breakdown so support engineers and DAG authors can
# attribute apparent slow startup to bundle prep (Airflow-side, e.g. git fetch) vs.
# DAG file parse (user code). Emitted before return so it lands in the task log.
# Under `run_as_user` impersonation, parse() runs once pre-sudo and again post-sudo,
# so this event fires twice for those tasks -- the two entries show impersonation overhead.
log.info(
"Worker startup parse complete",
bundle_name=bundle_info.name,
bundle_version=bundle_info.version,
dag_file=what.dag_rel_path,
dag_id=what.ti.dag_id,
bundle_prepare_ms=bundle_prepare_ms,
dag_file_parse_ms=dag_file_parse_ms,
)
return RuntimeTaskInstance.model_construct(
**what.ti.model_dump(exclude_unset=True),
task=task,
bundle_instance=bundle_instance,
_ti_context_from_server=what.ti_context,
max_tries=what.ti_context.max_tries,
start_date=what.start_date,
state=TaskInstanceState.RUNNING,
sentry_integration=what.sentry_integration,
)
# This global variable will be used by Connection/Variable/XCom classes, or other parts of the task's execution,
# to send requests back to the supervisor process.
#
# Why it needs to be a global:
# - Many parts of Airflow's codebase (e.g., connections, variables, and XComs) may rely on making dynamic requests
# to the parent process during task execution.
# - These calls occur in various locations and cannot easily pass the `CommsDecoder` instance through the
# deeply nested execution stack.
# - By defining `SUPERVISOR_COMMS` as a global, it ensures that this communication mechanism is readily
# accessible wherever needed during task execution without modifying every layer of the call stack.
SUPERVISOR_COMMS: CommsDecoder[ToTask, ToSupervisor]
# State machine!
# 1. Start up (receive details from supervisor)
# 2. Execution (run task code, possibly send requests)
# 3. Shutdown and report status
@detail_span("_verify_bundle_access")
def _verify_bundle_access(bundle_instance: BaseDagBundle, log: Logger) -> None:
"""
Verify bundle is accessible by the current user.
This is called after user impersonation (if any) to ensure the bundle
is actually accessible. Uses os.access() which works with any permission
scheme (standard Unix permissions, ACLs, SELinux, etc.).
:param bundle_instance: The bundle instance to check
:param log: Logger instance
:raises AirflowException: if bundle is not accessible
"""
from getpass import getuser
from airflow.sdk.exceptions import AirflowException
bundle_path = bundle_instance.path
if not bundle_path.exists():
# Already handled by initialize() with a warning
return
# Check read permission (and execute for directories to list contents)
access_mode = os.R_OK
if bundle_path.is_dir():
access_mode |= os.X_OK
if not os.access(bundle_path, access_mode):
raise AirflowException(
f"Bundle '{bundle_instance.name}' path '{bundle_path}' is not accessible "
f"by user '{getuser()}'. When using run_as_user, ensure bundle directories "
f"are readable by the impersonated user. "
f"See: https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/dag-bundles.html"
)
def get_startup_details() -> StartupDetails:
# The parent sends us a StartupDetails message un-prompted. After this, every single message is only sent
# in response to us sending a request.
if os.environ.get("_AIRFLOW__REEXECUTED_PROCESS") == "1" and (
msgjson := os.environ.get("_AIRFLOW__STARTUP_MSG")
):
# Clear any Kerberos replace cache if there is one, so new process can't reuse it.
os.environ.pop("KRB5CCNAME", None)
# entrypoint of re-exec process
msg: StartupDetails = TypeAdapter(StartupDetails).validate_json(msgjson)
reinit_supervisor_comms()
# We delay this message until _after_ we've got the logging re-configured, otherwise it will show up
# on stdout
log.debug("Using serialized startup message from environment", msg=msg)
else:
# normal entry point
msg = SUPERVISOR_COMMS._get_response() # type: ignore[assignment]
if not isinstance(msg, StartupDetails):
raise RuntimeError(f"Unhandled startup message {type(msg)} {msg}")
return msg
@detail_span("startup")
def startup(msg: StartupDetails) -> tuple[RuntimeTaskInstance, Context, Logger]:
# setproctitle causes issue on Mac OS: https://github.com/benoitc/gunicorn/issues/3021
os_type = sys.platform
if os_type == "darwin":
log.debug("Mac OS detected, skipping setproctitle")
else:
from setproctitle import setproctitle
setproctitle(f"airflow worker -- {msg.ti.id}")
# Bind TI identifiers so every subsequent log line from this worker process
# carries ti_id, enabling single-grep lifecycle reconstruction across components.
bind_contextvars(
ti_id=str(msg.ti.id),
dag_id=msg.ti.dag_id,
task_id=msg.ti.task_id,
run_id=msg.ti.run_id,
try_number=msg.ti.try_number,
map_index=msg.ti.map_index,
)
try:
get_listener_manager().hook.on_starting(component=TaskRunnerMarker())
except Exception:
log.exception("error calling listener")
with _airflow_parsing_context_manager(dag_id=msg.ti.dag_id, task_id=msg.ti.task_id):
ti = parse(msg, log)
log.debug("Dag file parsed", file=msg.dag_rel_path)
run_as_user = getattr(ti.task, "run_as_user", None) or conf.get(
"core", "default_impersonation", fallback=None
)
if os.environ.get("_AIRFLOW__REEXECUTED_PROCESS") != "1" and run_as_user and run_as_user != getuser():
# enters here for re-exec process
os.environ["_AIRFLOW__REEXECUTED_PROCESS"] = "1"
# store startup message in environment for re-exec process
os.environ["_AIRFLOW__STARTUP_MSG"] = msg.model_dump_json()
os.set_inheritable(SUPERVISOR_COMMS.socket.fileno(), True)
# Import main directly from the module instead of re-executing the file.
# This ensures that when other parts modules import
# airflow.sdk.execution_time.task_runner, they get the same module instance
# with the properly initialized SUPERVISOR_COMMS global variable.
# If we re-executed the module with `python -m`, it would load as __main__ and future
# imports would get a fresh copy without the initialized globals.
rexec_python_code = "from airflow.sdk.execution_time.task_runner import main; main()"
cmd = ["sudo", "-E", "-H", "-u", run_as_user, sys.executable, "-c", rexec_python_code]
log.info(
"Running command",
command=cmd,
)
os.execvp("sudo", cmd)
# ideally, we should never reach here, but if we do, we should return None, None, None
return None, None, None
return ti, ti.get_template_context(), log
def _serialize_template_field(
template_field: Any, name: str
) -> str | dict | list | int | float | bool | None:
"""
Return a serializable representation of the templated field.
The walk has two responsibilities:
1. **Make the template_field JSON-encodable** — every container is rebuilt
with primitive leaves (str/int/float/bool/None), tuples and sets are
flattened to lists, and unsupported objects fall through to ``str()``
so ``json.dumps`` never raises on the result.
2. **Keep the output deterministic across parses** — callables are replaced
with their qualified name (never the default ``<function ... at 0x...>``
repr), dicts are key-sorted, and (frozen)sets are sorted by element so
the same input always produces the same string.
Uses the SDK secrets masker to redact secrets in the serialized output.
"""
import inspect
from airflow.sdk._shared.module_loading import qualname
from airflow.sdk._shared.secrets_masker import redact
def normalize_dict_key(key) -> str:
"""Normalize a dict key to a serialized string type."""
# Serialized template_field keys must all be strings, not a mix of types, so that
# downstream json.dumps(..., sort_keys=True) does not raise on mixed-type keys.
return str(serialize_object(key))
def serialize_object(obj):
"""Recursively rewrite ``obj`` into a JSON-encodable, hash-stable structure."""
if obj is None or isinstance(obj, (str, int, float, bool)):
return obj
if isinstance(obj, dict):
# Serialize keys/values first so each key is a string and the output is hash-stable,
# then sort by the serialized key to prevent hash inconsistencies when dict ordering varies.
serialized_pairs = [(normalize_dict_key(k), serialize_object(v)) for k, v in obj.items()]
return dict(sorted(serialized_pairs, key=lambda kv: kv[0]))
if isinstance(obj, (list, tuple)):
return [serialize_object(item) for item in obj]
if isinstance(obj, (set, frozenset)):
# JSON has no set type → flatten to a list with deterministic ordering
# so hash randomization on element types cannot shift cross-process iteration order.
serialized_set = [serialize_object(e) for e in obj]
return sorted(serialized_set, key=lambda x: (type(x).__name__, str(x)))
# Use inspect.getattr_static to bypass any custom __getattr__ / metaclass magic
if callable(inspect.getattr_static(obj, "serialize", None)):
return serialize_object(obj.serialize())
# Kubernetes client objects (V1Pod, V1Container, ...) expose their content via to_dict().
# Scope the branch to the kubernetes namespace so unrelated user classes that happen to
# define a to_dict() method fall through to str() instead of being treated as K8s payloads.
if getattr(type(obj), "__module__", "").startswith(
("kubernetes.", "kubernetes_asyncio.")
) and callable(inspect.getattr_static(obj, "to_dict", None)):
return serialize_object(obj.to_dict())
if callable(obj):
# Use qualified name; default repr embeds memory addresses, which would change the DAG hash on every parse
return f"<callable {qualname(obj, True)}>"
# A custom __str__ or __repr__ is treated as an intentional textual representation
# supplied by the author and used as-is.
if type(obj).__str__ is not object.__str__ or type(obj).__repr__ is not object.__repr__:
return str(obj)
# Otherwise fall back to a qualname marker. The default object repr is
# `<ClassName object at 0x...>`, which embeds a memory address that flips per process
# and would break DAG hash stability — use the class qualname instead.
return f"<{qualname(type(obj), True)} object>"
max_length = conf.getint("core", "max_templated_field_length")
serialized = serialize_object(template_field)
if len(str(serialized)) > max_length:
# Redact while still structured to preserve nested-key context (so values under
# documented sensitive keys such as `password`, `token`, `secret`, `api_key`
# are masked recursively); only stringify the redacted result for truncation.
rendered = redact(serialized, name)
return truncate_rendered_value(str(rendered), max_length)
return serialized
@detail_span("_serialize_rendered_fields")
def _serialize_rendered_fields(task: AbstractOperator) -> dict[str, JsonValue]:
from airflow.sdk._shared.secrets_masker import redact
rendered_fields: dict[str, JsonValue] = {}
for field in task.template_fields:
value = getattr(task, field)
serialized = _serialize_template_field(value, field)
# Redact secrets in the task process itself before sending to API server
# This ensures that the secrets those are registered via mask_secret() on workers / dag processor are properly masked
# on the UI.
redacted = redact(serialized, field)
rendered_fields[field] = TypeAdapter(JsonValue).validate_python(redacted)
renderers = getattr(task, "template_fields_renderers", {})
for renderer_path in renderers:
if "." not in renderer_path:
continue
base_field, _, remainder = renderer_path.partition(".")
if base_field not in task.template_fields:
continue
base_value = getattr(task, base_field, None)
if base_value is None:
continue
current = base_value
for key in remainder.split("."):
if isinstance(current, dict):
current = current.get(key)
else:
current = getattr(current, key, None)
if current is None:
break
if current is not None:
serialized = _serialize_template_field(current, renderer_path)
redacted = redact(serialized, renderer_path)
rendered_fields[renderer_path] = TypeAdapter(JsonValue).validate_python(redacted)
return rendered_fields
def _build_asset_profiles(lineage_objects: list) -> Iterator[AssetProfile]:
# Lineage can have other types of objects besides assets, so we need to process them a bit.
for obj in lineage_objects or ():
if isinstance(obj, Asset):
yield AssetProfile(name=obj.name, uri=obj.uri, type=Asset.__name__)
elif isinstance(obj, AssetNameRef):
yield AssetProfile(name=obj.name, type=AssetNameRef.__name__)
elif isinstance(obj, AssetUriRef):
yield AssetProfile(uri=obj.uri, type=AssetUriRef.__name__)
elif isinstance(obj, AssetAlias):
yield AssetProfile(name=obj.name, type=AssetAlias.__name__)
def _serialize_outlet_events(events: OutletEventAccessorsProtocol) -> Iterator[dict[str, JsonValue]]:
if TYPE_CHECKING:
assert isinstance(events, OutletEventAccessors)
# We just collect everything the user recorded in the accessors.
# Further filtering will be done in the API server.
for key, accessor in events._dict.items():
if isinstance(key, AssetUniqueKey):
if accessor.partition_keys:
for partition_key in accessor.partition_keys:
yield {
"dest_asset_key": attrs.asdict(key),
"extra": accessor.extra,
"partition_key": partition_key,
}
else:
yield {"dest_asset_key": attrs.asdict(key), "extra": accessor.extra}
for alias_event in accessor.asset_alias_events:
yield attrs.asdict(alias_event)
@detail_span("_prepare")
def _prepare(ti: RuntimeTaskInstance, log: Logger, context: Context) -> ToSupervisor | None:
ti.hostname = get_hostname()
ti.task = ti.task.prepare_for_execution()
# Since context is now cached, and calling `ti.get_template_context` will return the same dict, we want to
# update the value of the task that is sent from there
context["task"] = ti.task
jinja_env = ti.task.dag.get_template_env()
ti.render_templates(context=context, jinja_env=jinja_env)
if rendered_fields := _serialize_rendered_fields(ti.task):
# so that we do not call the API unnecessarily
SUPERVISOR_COMMS.send(msg=SetRenderedFields(rendered_fields=rendered_fields))
# Try to render map_index_template early with available context (will be re-rendered after execution)
# This provides a partial label during task execution for templates using pre-execution context
# If rendering fails here, we suppress the error since it will be re-rendered after execution
try:
if rendered_map_index := _render_map_index(context, ti=ti, log=log):
ti.rendered_map_index = rendered_map_index
log.debug("Sending early rendered map index", length=len(rendered_map_index))
SUPERVISOR_COMMS.send(msg=SetRenderedMapIndex(rendered_map_index=rendered_map_index))
except Exception:
log.debug(
"Early rendering of map_index_template failed, will retry after task execution", exc_info=True
)
_validate_task_inlets_and_outlets(ti=ti, log=log)
try:
# TODO: Call pre execute etc.
get_listener_manager().hook.on_task_instance_running(
previous_state=TaskInstanceState.QUEUED, task_instance=ti
)
except Exception:
log.exception("error calling listener")
# No error, carry on and execute the task
return None
@detail_span("_validate_task_inlets_and_outlets")
def _validate_task_inlets_and_outlets(*, ti: RuntimeTaskInstance, log: Logger) -> None:
if not ti.task.inlets and not ti.task.outlets:
return
inactive_assets_resp = SUPERVISOR_COMMS.send(msg=ValidateInletsAndOutlets(ti_id=ti.id))
if TYPE_CHECKING:
assert isinstance(inactive_assets_resp, InactiveAssetsResult)
if inactive_assets := inactive_assets_resp.inactive_assets:
raise AirflowInactiveAssetInInletOrOutletException(
inactive_asset_keys=[
AssetUniqueKey.from_profile(asset_profile) for asset_profile in inactive_assets
]
)
def _defer_task(
defer: TaskDeferred, ti: RuntimeTaskInstance, log: Logger
) -> tuple[ToSupervisor, TaskInstanceState]:
log.info("Pausing task as DEFERRED. ", dag_id=ti.dag_id, task_id=ti.task_id, run_id=ti.run_id)
classpath, trigger_kwargs = defer.trigger.serialize()
queue: str | None = None
# Only inherit the deferring task's queue when triggerer.queues_enabled conf is True
# and the trigger doesn't already manage its own queue (e.g. BaseEventTrigger, CallbackTrigger).
if conf.getboolean("triggerer", "queues_enabled", fallback=False) and getattr(
defer.trigger, "trigger_queue_inherited_from_task", True
):
queue = ti.task.queue
from airflow.sdk.serde import serialize as serde_serialize
trigger_kwargs = serde_serialize(trigger_kwargs)
next_kwargs = serde_serialize(defer.kwargs or {})
if TYPE_CHECKING:
assert isinstance(next_kwargs, dict)
assert isinstance(trigger_kwargs, dict)
msg = DeferTask(
classpath=classpath,
trigger_kwargs=trigger_kwargs,
trigger_timeout=defer.timeout,
queue=queue,
next_method=defer.method_name,
next_kwargs=next_kwargs,
)
state = TaskInstanceState.DEFERRED
return msg, state
def _await_input_task(
awaiting: TaskAwaitingInput, ti: RuntimeTaskInstance, log: Logger
) -> tuple[ToSupervisor, TaskInstanceState]:
"""Build the message that parks a task in AWAITING_INPUT (HITL), with no trigger."""
log.info("Pausing task as AWAITING_INPUT.", dag_id=ti.dag_id, task_id=ti.task_id, run_id=ti.run_id)
from airflow.sdk.serde import serialize as serde_serialize
next_kwargs = serde_serialize(awaiting.kwargs or {})
if TYPE_CHECKING:
assert isinstance(next_kwargs, dict)
msg = AwaitInputTask(
timeout=awaiting.timeout,
next_method=awaiting.method_name,
next_kwargs=next_kwargs,
)
state = TaskInstanceState.AWAITING_INPUT
return msg, state
def _run_task_and_map_outcome(
ti: RuntimeTaskInstance,
context: Context,
log: Logger,
) -> tuple[TaskInstanceState, ToSupervisor | None, BaseException | None]:
"""Execute the task and map its outcome -- success or a handled exception -- to a message and state."""
import signal
from airflow.sdk.exceptions import (
AirflowRescheduleException,
AirflowSkipException,
DagRunTriggerException,
DownstreamTasksSkipped,
TaskAwaitingInput,
TaskDeferred,
)
if TYPE_CHECKING:
assert ti.task is not None
assert isinstance(ti.task, BaseOperator)
parent_pid = os.getpid()
def _on_term(signum, frame):
pid = os.getpid()
if pid != parent_pid:
return
ti.task.on_kill()
signal.signal(signal.SIGTERM, _on_term)
msg: ToSupervisor | None = None
state: TaskInstanceState | None = None
error: BaseException | None = None
try:
# First, clear the xcom data sent from server
if ti._ti_context_from_server and (keys_to_delete := ti._ti_context_from_server.xcom_keys_to_clear):
for x in keys_to_delete:
log.debug("Clearing XCom with key", key=x)
XCom.delete(
key=x,
dag_id=ti.dag_id,
task_id=ti.task_id,
run_id=ti.run_id,
map_index=ti.map_index,
)
with set_current_context(context):
# This is the earliest that we can render templates -- as if it excepts for any reason we need to
# catch it and handle it like a normal task failure
if early_exit := _prepare(ti, log, context):
msg = early_exit
ti.state = state = TaskInstanceState.FAILED
return state, msg, error
try:
result = _execute_task(context=context, ti=ti, log=log)
log.info("::group::Post Execute")
except Exception:
import jinja2
# If the task failed, swallow rendering error so it doesn't mask the main error.
with suppress(jinja2.TemplateSyntaxError, jinja2.UndefinedError):
previous_rendered_map_index = ti.rendered_map_index
ti.rendered_map_index = _render_map_index(context, ti=ti, log=log)
# Send update only if value changed (e.g., user set context variables during execution)
if ti.rendered_map_index and ti.rendered_map_index != previous_rendered_map_index:
SUPERVISOR_COMMS.send(
msg=SetRenderedMapIndex(rendered_map_index=ti.rendered_map_index)
)
raise
else: # If the task succeeded, render normally to let rendering error bubble up.
previous_rendered_map_index = ti.rendered_map_index
ti.rendered_map_index = _render_map_index(context, ti=ti, log=log)
# Send update only if value changed (e.g., user set context variables during execution)
if ti.rendered_map_index and ti.rendered_map_index != previous_rendered_map_index:
SUPERVISOR_COMMS.send(msg=SetRenderedMapIndex(rendered_map_index=ti.rendered_map_index))
_push_xcom_if_needed(result, ti, log)
msg, state = _handle_current_task_success(context, ti)
except DownstreamTasksSkipped as skip:
log.info("::group::Post Execute")
log.info("Skipping downstream tasks.")
tasks_to_skip = skip.tasks if isinstance(skip.tasks, list) else [skip.tasks]
SUPERVISOR_COMMS.send(msg=SkipDownstreamTasks(tasks=tasks_to_skip))
msg, state = _handle_current_task_success(context, ti)
except DagRunTriggerException as drte:
log.info("::group::Post Execute")
msg, state = _handle_trigger_dag_run(drte, context, ti, log)
except TaskDeferred as defer:
log.info("::group::Post Execute")
msg, state = _defer_task(defer, ti, log)
except TaskAwaitingInput as awaiting:
log.info("::group::Post Execute")
msg, state = _await_input_task(awaiting, ti, log)
except AirflowSkipException as e:
log.info("::group::Post Execute")
if e.args:
log.info("Skipping task.", reason=e.args[0])
msg = TaskState(
state=TaskInstanceState.SKIPPED,
end_date=datetime.now(tz=timezone.utc),
rendered_map_index=ti.rendered_map_index,
)
state = TaskInstanceState.SKIPPED
except AirflowRescheduleException as reschedule:
log.info("::group::Post Execute")
log.info("Rescheduling task, marking task as UP_FOR_RESCHEDULE")
msg = RescheduleTask(
reschedule_date=reschedule.reschedule_date, end_date=datetime.now(tz=timezone.utc)
)
state = TaskInstanceState.UP_FOR_RESCHEDULE
except (AirflowFailException, AirflowSensorTimeout) as e:
# If AirflowFailException is raised, task should not retry.
# If a sensor in reschedule mode reaches timeout, task should not retry.
log.exception("Task failed with exception")
log.info("::group::Post Execute")
ti.end_date = datetime.now(tz=timezone.utc)
msg = TaskState(
state=TaskInstanceState.FAILED,
end_date=ti.end_date,
rendered_map_index=ti.rendered_map_index,
)
state = TaskInstanceState.FAILED
error = e
except (AirflowTaskTimeout, AirflowException, AirflowRuntimeError) as e:
# We should allow retries if the task has defined it.
log.exception("Task failed with exception")
log.info("::group::Post Execute")
msg, state = _handle_current_task_failed(ti, e, log, context)
error = e
except AirflowTaskTerminated as e:
# External state updates are already handled with `ti_heartbeat` and will be
# updated already be another UI API. So, these exceptions should ideally never be thrown.
# If these are thrown, we should mark the TI state as failed.
log.exception("Task failed with exception")
log.info("::group::Post Execute")
ti.end_date = datetime.now(tz=timezone.utc)
msg = TaskState(
state=TaskInstanceState.FAILED,
end_date=ti.end_date,
rendered_map_index=ti.rendered_map_index,
)
state = TaskInstanceState.FAILED
error = e
except SystemExit as e:
# SystemExit needs to be retried if they are eligible.
log.error("Task exited", exit_code=e.code)
log.info("::group::Post Execute")
msg, state = _handle_current_task_failed(ti, e, log, context)
error = e
except BaseException as e:
log.exception("Task failed with exception")
log.info("::group::Post Execute")
msg, state = _handle_current_task_failed(ti, e, log, context)
error = e
return state, msg, error
@Sentry.enrich_errors
@detail_span("run")
def run(
ti: RuntimeTaskInstance,
context: Context,
log: Logger,
) -> tuple[TaskInstanceState, ToSupervisor | None, BaseException | None]:
"""Run the task in this process."""
msg: ToSupervisor | None = None
state: TaskInstanceState | None = None
error: BaseException | None = None
stats_tags = ti.stats_tags
stats.incr("ti.start", tags=stats_tags)
try:
state, msg, error = _run_task_and_map_outcome(ti, context, log)
except Exception as e:
# Python does not offer an exception raised inside an ``except`` clause to that
# clause's siblings, so a handler that fails would otherwise escape entirely --
# skipping the retry decision and every callback in ``finalize()``.
msg, state, error = _handle_handler_failure(ti, e, log, context)
finally:
# `state` may still be unset if an exception handler above raised before
# binding it
if state is not None:
stats.incr("ti.finish", tags={**stats_tags, "state": state.value})
if msg:
# If the supervisor rejects the terminal-state report
# (e.g. the server already moved the TI to a terminal state and
# returns 409 - for example due to issue fixed by #65594),
# we still want to run finalize() so listeners observe the task state.
try:
SUPERVISOR_COMMS.send(msg=msg)
except Exception:
log.exception(
"Failed to report terminal task state to supervisor",
state=state.value if state is not None else None,
)
# Fail closed for FAILED / UP_FOR_RETRY: when the supervisor
# never receives the terminal-state message, exiting 0 would
# let the supervisor's final_state property default to
# SUCCESS (exit_code == 0 with no _terminal_state set),
# turning a real failure into a silent data-quality bug for
# every downstream task. We signal main() to sys.exit(1)
# AFTER finalize() runs, so on_failure_callback /
# on_retry_callback / listener hooks / email_on_failure /
# email_on_retry still fire. sys.exit(1) directly here would
# raise SystemExit, which is BaseException, not Exception —
# main()'s `except Exception:` would not catch it and
# finalize() at the call site would be skipped.
#
# SKIPPED / UP_FOR_RESCHEDULE / DEFERRED are intentionally
# not fail-closed: supervisor's final_state would misclassify
# them too, but exiting non-zero would map them to FAILED,
# which is strictly worse than the default. Those need a
# separate fix in supervisor's final_state.
if state in (TaskInstanceState.FAILED, TaskInstanceState.UP_FOR_RETRY):
ti._terminal_state_send_failed = True
# Return the message to make unit tests easier too
ti.state = state
return state, msg, error
def _handle_current_task_success(
context: Context,
ti: RuntimeTaskInstance,
) -> tuple[SucceedTask, TaskInstanceState]:
end_date = datetime.now(tz=timezone.utc)
ti.end_date = end_date
# Record operator and task instance success metrics
operator = ti.task.__class__.__name__
stats_tags = ti.stats_tags
stats.incr("operator_successes", tags={**stats_tags, "operator_name": operator})
stats.incr("ti_successes", tags=stats_tags)
task_outlets = list(_build_asset_profiles(ti.task.outlets))
outlet_events = list(_serialize_outlet_events(context["outlet_events"]))
if conf.getboolean("state_store", "clear_on_success") and _get_worker_state_store_backend() is not None:
log.info(
"Clearing task state from custom backend as clear_on_success is enabled. The database references will be cleared by the API server."
)
context["task_state_store"]._clear_backend_only()
msg = SucceedTask(
end_date=end_date,
task_outlets=task_outlets,
outlet_events=outlet_events,
rendered_map_index=ti.rendered_map_index,
)
return msg, TaskInstanceState.SUCCESS
def _evaluate_retry_policy(
ti: RuntimeTaskInstance,
exception: BaseException,
log: Logger,
context: Context | None = None,
) -> RetryDecision | None:
"""
Evaluate the task's retry policy if one is configured.
Returns ``None`` when no policy is configured so the caller falls through
to the standard retry logic.
"""
policy = getattr(ti.task, "retry_policy", None)
if policy is None:
return None
try:
max_tries = ti._ti_context_from_server.max_tries if ti._ti_context_from_server else 0
decision = policy.evaluate(
exception=exception,
try_number=ti.try_number,
max_tries=max_tries,
context=context,
)
if decision.reason:
# Close the group so the retry policy decision is not hidden inside "Post Execute".
log.info("::endgroup::")
log.info("Retry policy decision", action=decision.action.value, reason=decision.reason)
return decision
except Exception:
log.exception("Retry policy evaluation failed, falling back to default behaviour")
return None
def _handle_handler_failure(
ti: RuntimeTaskInstance, exception: Exception, log: Logger, context: Context
) -> tuple[RetryTask | TaskState, TaskInstanceState, Exception]:
"""
Decide the outcome for an exception raised by one of the outcome handlers themselves.
Handlers reach this by talking to the API server -- a missing Dag makes
``_handle_trigger_dag_run`` raise ``AirflowRuntimeError`` -- or by serializing
user-supplied values, since ``_defer_task`` and ``_await_input_task`` run
``serde_serialize`` over kwargs and it raises ``TypeError`` for anything it has no
serializer for.
The handler chain's own classifications are preserved rather than routing everything
through the retry-count check, so an exception that means "do not retry" still means
that when it surfaces from a handler.
"""
log.exception("Task failed with exception")
if isinstance(exception, (AirflowFailException, AirflowSensorTimeout, AirflowTaskTerminated)):
return _terminal_failure(ti), TaskInstanceState.FAILED, exception
try:
msg, state = _handle_current_task_failed(ti, exception, log, context)
except Exception:
# The failure path itself failed. Re-entering it would just fail again and
# escape, losing the callbacks this whole function exists to preserve, so fail
# closed on a plain FAILED state instead.
log.exception("Could not determine terminal state, failing closed")
return _terminal_failure(ti), TaskInstanceState.FAILED, exception
return msg, state, exception
def _terminal_failure(ti: RuntimeTaskInstance) -> TaskState:
"""Build a plain FAILED terminal message, bypassing any retry decision."""
ti.end_date = datetime.now(tz=timezone.utc)
return TaskState(
state=TaskInstanceState.FAILED,
end_date=ti.end_date,
rendered_map_index=ti.rendered_map_index,
)
def _handle_current_task_failed(
ti: RuntimeTaskInstance,
exception: BaseException,
log: Logger,
context: Context | None = None,
) -> tuple[RetryTask | TaskState, TaskInstanceState]:
"""
Handle a failed task: evaluate the retry policy (if any) and decide the next state.
This is the entry point every failure path routes through. When the policy
returns FAIL the task is marked as failed immediately, bypassing the normal
retry-count check. When it returns RETRY with a custom delay, that delay is
forwarded in the ``RetryTask`` message. For DEFAULT (or when no policy is
configured), the standard ``_finalize_task_failure`` logic runs.
"""
from airflow.sdk.definitions.retry_policy import RetryAction
decision = _evaluate_retry_policy(ti, exception, log, context)
if decision is not None and decision.action == RetryAction.FAIL:
ti.end_date = datetime.now(tz=timezone.utc)
return (
TaskState(
state=TaskInstanceState.FAILED,
end_date=ti.end_date,
rendered_map_index=ti.rendered_map_index,
),
TaskInstanceState.FAILED,
)
if decision is not None and decision.action == RetryAction.RETRY:
return _finalize_task_failure(
ti, retry_delay_override=decision.retry_delay, retry_reason=decision.reason
)
return _finalize_task_failure(ti)
def _finalize_task_failure(
ti: RuntimeTaskInstance,
retry_delay_override: timedelta | None = None,
retry_reason: str | None = None,
) -> tuple[RetryTask, TaskInstanceState] | tuple[TaskState, TaskInstanceState]:
"""
Record failure metrics and build the standard retry-or-fail outcome.
Returns an ``UP_FOR_RETRY`` ``RetryTask`` when the server marked the task
retry-eligible (optionally carrying a policy-supplied delay/reason), else a
``FAILED`` ``TaskState``. This is the default path; retry-policy overrides
are decided in :func:`_handle_current_task_failed` before this is called.
"""
end_date = datetime.now(tz=timezone.utc)
ti.end_date = end_date
# Record operator and task instance failed metrics. One failure is one increment even
# if this runs twice, which happens when a first pass raised after counting -- see
# `_handle_handler_failure`.
if not ti._failure_metrics_emitted:
operator = ti.task.__class__.__name__
stats_tags = ti.stats_tags
stats.incr("operator_failures", tags={**stats_tags, "operator_name": operator})
stats.incr("ti_failures", tags=stats_tags)
ti._failure_metrics_emitted = True
if ti._ti_context_from_server and ti._ti_context_from_server.should_retry:
retry_kwargs: dict[str, Any] = {"end_date": end_date}
if retry_delay_override is not None:
retry_kwargs["retry_delay_seconds"] = retry_delay_override.total_seconds()
if retry_reason is not None:
retry_kwargs["retry_reason"] = retry_reason[:500]
return RetryTask(**retry_kwargs), TaskInstanceState.UP_FOR_RETRY
return (
TaskState(
state=TaskInstanceState.FAILED, end_date=end_date, rendered_map_index=ti.rendered_map_index
),
TaskInstanceState.FAILED,
)
def _handle_trigger_dag_run(
drte: DagRunTriggerException, context: Context, ti: RuntimeTaskInstance, log: Logger
) -> tuple[ToSupervisor, TaskInstanceState]:
"""Handle exception from TriggerDagRunOperator."""
log.info("Triggering Dag Run.", trigger_dag_id=drte.trigger_dag_id)
comms_msg = SUPERVISOR_COMMS.send(
TriggerDagRun(
dag_id=drte.trigger_dag_id,
run_id=drte.dag_run_id,
logical_date=drte.logical_date,
run_after=drte.run_after,
conf=drte.conf,
reset_dag_run=drte.reset_dag_run,
note=drte.note,
),
)
if isinstance(comms_msg, ErrorResponse) and comms_msg.error == ErrorType.DAGRUN_ALREADY_EXISTS:
if drte.skip_when_already_exists:
log.info(
"Dag Run already exists, skipping task as skip_when_already_exists is set to True.",
dag_id=drte.trigger_dag_id,
)
msg = TaskState(
state=TaskInstanceState.SKIPPED,
end_date=datetime.now(tz=timezone.utc),
rendered_map_index=ti.rendered_map_index,
)
state = TaskInstanceState.SKIPPED
else:
log.error("Dag Run already exists, marking task as failed.", dag_id=drte.trigger_dag_id)
msg = TaskState(
state=TaskInstanceState.FAILED,
end_date=datetime.now(tz=timezone.utc),
rendered_map_index=ti.rendered_map_index,
)
state = TaskInstanceState.FAILED
return msg, state
log.info("Dag Run triggered successfully.", trigger_dag_id=drte.trigger_dag_id)
# Store the run id from the dag run (either created or found above) to
# be used when creating the extra link on the webserver.
ti.xcom_push(key="trigger_run_id", value=drte.dag_run_id)
if drte.wait_for_completion:
if drte.deferrable:
from airflow.providers.standard.triggers.external_task import DagStateTrigger
defer = TaskDeferred(
trigger=DagStateTrigger(
dag_id=drte.trigger_dag_id,
states=drte.allowed_states + drte.failed_states, # type: ignore[arg-type]
# Don't filter by execution_dates when run_ids is provided.
# run_id uniquely identifies a DAG run, and when reset_dag_run=True,
# drte.logical_date might be a newly calculated value that doesn't match
# the persisted logical_date in the database, causing the trigger to never find the run.
execution_dates=None,
run_ids=[drte.dag_run_id],
poll_interval=drte.poke_interval,
),
method_name="execute_complete",
)
return _defer_task(defer, ti, log)
while True:
log.info(
"Waiting for dag run to complete execution in allowed state.",
dag_id=drte.trigger_dag_id,
run_id=drte.dag_run_id,
allowed_state=drte.allowed_states,
)
time.sleep(drte.poke_interval)
comms_msg = SUPERVISOR_COMMS.send(
GetDagRunState(dag_id=drte.trigger_dag_id, run_id=drte.dag_run_id)
)
if TYPE_CHECKING:
assert isinstance(comms_msg, DagRunStateResult)
if comms_msg.state in drte.failed_states:
log.error(
"DagRun finished with failed state.", dag_id=drte.trigger_dag_id, state=comms_msg.state
)
# Mirror the deferrable path (DagStateTrigger -> execute_complete raises
# AirflowException), which flows through run()'s exception handler and
# therefore honours a configured retry_policy. Synthesize the same
# exception here so non-deferrable waits evaluate the policy too,
# falling back to the standard retry-count check when none is set.
return _handle_current_task_failed(
ti,
AirflowException(f"{drte.trigger_dag_id} failed with failed state {comms_msg.state}"),
log,
context,
)
if comms_msg.state in drte.allowed_states:
log.info(
"DagRun finished with allowed state.", dag_id=drte.trigger_dag_id, state=comms_msg.state
)
break
log.debug(
"DagRun not yet in allowed or failed state.",
dag_id=drte.trigger_dag_id,
state=comms_msg.state,
)
else:
# Fire-and-forget mode: wait_for_completion=False
if drte.deferrable:
log.info(
"Ignoring deferrable=True because wait_for_completion=False. "
"Task will complete immediately without waiting for the triggered DAG run.",
trigger_dag_id=drte.trigger_dag_id,
)
return _handle_current_task_success(context, ti)
def _run_task_state_change_callbacks(
task: BaseOperator,
kind: Literal[
"on_execute_callback",
"on_failure_callback",
"on_success_callback",
"on_retry_callback",
"on_skipped_callback",
],
context: Context,
log: Logger,
) -> None:
callback: Callable[[Context], None]
for i, callback in enumerate(getattr(task, kind)):
try:
create_executable_runner(callback, context_get_outlet_events(context), logger=log).run(context)
except Exception:
log.exception("Failed to run task callback", kind=kind, index=i, callback=callback)
def _send_error_email_notification(
task: BaseOperator | MappedOperator,
ti: RuntimeTaskInstance,
context: Context,
error: BaseException | str | None,
log: Logger,
) -> None:
"""
Send email notification for task errors through the configured email backend.
A non-default ``[email] email_backend`` (an SES, SendGrid or org-internal callable with the
``airflow.utils.email.send_email`` signature) is wrapped in
:class:`~airflow.sdk.execution_time.email_backend._LegacyEmailBackendNotifier`; otherwise the
default :class:`~airflow.providers.smtp.notifications.smtp.SmtpNotifier` is used.
Both the worker task-runner path (:func:`finalize`) and the DAG-processor callback path
(``_execute_email_callbacks``) funnel through this function, so the resolved backend is used
consistently regardless of how the task failed.
"""
if not task.email:
return
email_backend = conf.get("email", "email_backend", fallback=_DEFAULT_EMAIL_BACKEND)
notifier_description = "SmtpNotifier"
if email_backend and email_backend != _DEFAULT_EMAIL_BACKEND:
notifier_class: _ErrorEmailNotifier = _LegacyEmailBackendNotifier
notifier_description = f"configured email_backend {email_backend!r}"
else:
try:
from airflow.providers.smtp.notifications.smtp import SmtpNotifier
except ImportError:
log.error(
"Failed to send task failure or retry email notification: "
"`apache-airflow-providers-smtp` is not installed. "
"Install this provider to enable email notifications."
)
return
notifier_class = SmtpNotifier
subject_template_file = conf.get("email", "subject_template", fallback=None)
# Read the template file if configured
if subject_template_file and Path(subject_template_file).exists():
subject = Path(subject_template_file).read_text()
else:
# Fallback to default
subject = "Airflow alert: {{ti}}"
html_content_template_file = conf.get("email", "html_content_template", fallback=None)
# Read the template file if configured
if html_content_template_file and Path(html_content_template_file).exists():
html_content = Path(html_content_template_file).read_text()
else:
# Fallback to default
# For reporting purposes, we report based on 1-indexed,
# not 0-indexed lists (i.e. Try 1 instead of Try 0 for the first attempt).
html_content = (
"Try {{try_number}} out of {{max_tries + 1}}<br>"
"Exception:<br>{{exception_html}}<br>"
'Log: <a href="{{ti.log_url}}">Link</a><br>'
"Host: {{ti.hostname}}<br>"
'Mark success: <a href="{{ti.mark_success_url}}">Link</a><br>'
)
# Add exception_html to context for template rendering
import html
exception_html = html.escape(str(error)).replace("\n", "<br>")
additional_context = {
"exception": error,
"exception_html": exception_html,
"try_number": ti.try_number,
"max_tries": ti.max_tries,
}
email_context = {**context, **additional_context}
to_emails = task.email
if not to_emails:
return
try:
notifier = notifier_class(
to=to_emails,
subject=subject,
html_content=html_content,
from_email=conf.get("email", "from_email", fallback="airflow@airflow"),
)
notifier(email_context)
except Exception:
log.exception("Failed to send email notification via %s", notifier_description)
@detail_span("task.execute")
def _run_execute_callable(
context: Context,
execute: Callable[..., Any] | functools.partial[Any],
task: BaseOperator,
) -> Any:
"""
Run the task's execute callable, applying the execution timeout if one is set.
The contextvars snapshot is taken here, after the ``task.execute`` span is
current, so spans the operator emits during ``execute`` nest under it rather
than under the caller. ``ExecutorSafeguard``'s tracker is set into that copy
so the operator's ``execute`` passes the safeguard check, while the copy keeps
the change from leaking into the surrounding context.
"""
ctx = contextvars.copy_context()
ctx.run(ExecutorSafeguard.tracker.set, task)
if task.execution_timeout:
from airflow.sdk.execution_time.timeout import timeout
# TODO: handle timeout in case of deferral
timeout_seconds = task.execution_timeout.total_seconds()
try:
# It's possible we're already timed out, so fast-fail if true
if timeout_seconds <= 0:
raise AirflowTaskTimeout()
# Run task in timeout wrapper
with timeout(timeout_seconds):
result = ctx.run(execute, context=context)
except AirflowTaskTimeout:
task.on_kill()
raise
else:
result = ctx.run(execute, context=context)
return result
@detail_span("_execute_task")
def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger):
"""Execute Task (optionally with a Timeout) and push Xcom results."""
task = ti.task
execute = task.execute
if ti._ti_context_from_server and (next_method := ti._ti_context_from_server.next_method):
from airflow.sdk.serde import deserialize
next_kwargs_data = ti._ti_context_from_server.next_kwargs or {}
try:
if TYPE_CHECKING:
assert isinstance(next_kwargs_data, dict)
kwargs = deserialize(next_kwargs_data)
except (ImportError, KeyError, AttributeError, TypeError):
from airflow.serialization.serialized_objects import BaseSerialization
kwargs = BaseSerialization.deserialize(next_kwargs_data)
if TYPE_CHECKING:
assert isinstance(kwargs, dict)
execute = functools.partial(task.resume_execution, next_method=next_method, next_kwargs=kwargs)
# Export context in os.environ to make it available for operators to use.
airflow_context_vars = context_to_airflow_vars(context, in_env_var_format=True)
os.environ.update(airflow_context_vars)
outlet_events = context_get_outlet_events(context)
if (pre_execute_hook := task._pre_execute_hook) is not None:
create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context)
if getattr(pre_execute_hook := task.pre_execute, "__func__", None) is not BaseOperator.pre_execute:
create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context)
_run_task_state_change_callbacks(task, "on_execute_callback", context, log)
log.info("::endgroup::")
result = _run_execute_callable(context, execute, task)
if (post_execute_hook := task._post_execute_hook) is not None:
create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context, result)
if getattr(post_execute_hook := task.post_execute, "__func__", None) is not BaseOperator.post_execute:
create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context)
return result
def _render_map_index(context: Context, ti: RuntimeTaskInstance, log: Logger) -> str | None:
"""Render named map index if the Dag author defined map_index_template at the task level."""
if (template := context.get("map_index_template")) is None:
return None
log.debug("Rendering map_index_template", template_length=len(template))
jinja_env = ti.task.dag.get_template_env()
rendered_map_index = jinja_env.from_string(template).render(context)
log.debug("Map index rendered", length=len(rendered_map_index))
return rendered_map_index
def _push_xcom_if_needed(result: Any, ti: RuntimeTaskInstance, log: Logger):
"""Push XCom values when task has ``do_xcom_push`` set to ``True`` and the task returns a result."""
if ti.task.do_xcom_push:
xcom_value = result
else:
xcom_value = None
has_mapped_dep = next(ti.task.iter_mapped_dependants(), None) is not None
if xcom_value is None:
if not ti.is_mapped and has_mapped_dep:
# Uhoh, a downstream mapped task depends on us to push something to map over
from airflow.sdk.exceptions import XComForMappingNotPushed
raise XComForMappingNotPushed()
return
mapped_length: int | None = None
if not ti.is_mapped and has_mapped_dep:
from airflow.sdk.definitions.mappedoperator import is_mappable_value
from airflow.sdk.exceptions import UnmappableXComTypePushed
if not is_mappable_value(xcom_value):
raise UnmappableXComTypePushed(xcom_value)
mapped_length = len(xcom_value)
log.info("Pushing xcom", ti=ti)
# If the task has multiple outputs, push each output as a separate XCom.
if ti.task.multiple_outputs:
if not isinstance(xcom_value, Mapping):
raise TypeError(
f"Returned output was type {type(xcom_value)} expected dictionary for multiple_outputs"
)
for key in xcom_value.keys():
if not isinstance(key, str):
raise TypeError(
"Returned dictionary keys must be strings when using "
f"multiple_outputs, found {key} ({type(key)}) instead"
)
for k, v in result.items():
ti.xcom_push(k, v)
_xcom_push(ti, BaseXCom.XCOM_RETURN_KEY, result, mapped_length=mapped_length)
@detail_span("finalize")
def finalize(
ti: RuntimeTaskInstance,
state: TaskInstanceState,
context: Context,
log: Logger,
error: BaseException | None = None,
):
# Record task duration metrics for all terminal states
if ti.start_date and ti.end_date:
duration_ms = (ti.end_date - ti.start_date).total_seconds() * 1000
stats.timing("task.duration", duration_ms, tags=ti.stats_tags)
task = ti.task
# Pushing xcom for each operator extra links defined on the operator only.
for oe in task.operator_extra_links:
try:
link, xcom_key = oe.get_link(operator=task, ti_key=ti), oe.xcom_key # type: ignore[arg-type]
log.debug("Setting xcom for operator extra link", link=link, xcom_key=xcom_key)
_xcom_push_to_db(ti, key=xcom_key, value=link)
except Exception:
log.exception(
"Failed to push an xcom for task operator extra link",
link_name=oe.name,
xcom_key=oe.xcom_key,
ti=ti,
)
if getattr(ti.task, "overwrite_rtif_after_execution", False):
log.debug("Overwriting Rendered template fields.")
if ti.task.template_fields:
try:
SUPERVISOR_COMMS.send(SetRenderedFields(rendered_fields=_serialize_rendered_fields(ti.task)))
except Exception:
log.exception("Failed to set rendered fields during finalization", ti=ti, task=ti.task)
log.debug("Running finalizers", ti=ti)
if state == TaskInstanceState.SUCCESS:
_run_task_state_change_callbacks(task, "on_success_callback", context, log)
try:
get_listener_manager().hook.on_task_instance_success(
previous_state=TaskInstanceState.RUNNING, task_instance=ti
)
except Exception:
log.exception("error calling listener")
elif state == TaskInstanceState.SKIPPED:
_run_task_state_change_callbacks(task, "on_skipped_callback", context, log)
try:
get_listener_manager().hook.on_task_instance_skipped(
previous_state=TaskInstanceState.RUNNING, task_instance=ti
)
except Exception:
log.exception("error calling listener")
elif state == TaskInstanceState.UP_FOR_RETRY:
_run_task_state_change_callbacks(task, "on_retry_callback", context, log)
try:
get_listener_manager().hook.on_task_instance_failed(
previous_state=TaskInstanceState.RUNNING, task_instance=ti, error=error
)
except Exception:
log.exception("error calling listener")
if error and task.email_on_retry and task.email:
_send_error_email_notification(task, ti, context, error, log)
elif state == TaskInstanceState.FAILED:
_run_task_state_change_callbacks(task, "on_failure_callback", context, log)
try:
get_listener_manager().hook.on_task_instance_failed(
previous_state=TaskInstanceState.RUNNING, task_instance=ti, error=error
)
except Exception:
log.exception("error calling listener")
if error and task.email_on_failure and task.email:
_send_error_email_notification(task, ti, context, error, log)
try:
get_listener_manager().hook.before_stopping(component=TaskRunnerMarker())
except Exception:
log.exception("error calling listener")
log.info("::endgroup::")
@contextmanager
def flush_spans():
try:
yield
finally:
provider = trace.get_tracer_provider()
if hasattr(provider, "force_flush"):
from airflow.sdk.configuration import conf
timeout_millis = conf.getint("traces", "task_runner_flush_timeout_milliseconds", fallback=30000)
provider.force_flush(timeout_millis=timeout_millis)
@flush_spans()
def main():
log = structlog.get_logger(logger_name="task")
global SUPERVISOR_COMMS
SUPERVISOR_COMMS = CommsDecoder[ToTask, ToSupervisor](log=log)
stats.initialize(
factory=stats_utils.get_stats_factory(),
export_legacy_names=conf.getboolean("metrics", "legacy_names_on"),
)
stack = ExitStack()
span = INVALID_SPAN
with stack:
try:
try:
log.info("::group::Pre Execute")
startup_details = get_startup_details()
span_ctx_mgr = _make_task_span(msg=startup_details)
span = stack.enter_context(span_ctx_mgr)
ti, context, log = startup(msg=startup_details)
except AirflowRescheduleException as reschedule:
log.warning("Rescheduling task during startup, marking task as UP_FOR_RESCHEDULE")
SUPERVISOR_COMMS.send(
msg=RescheduleTask(
reschedule_date=reschedule.reschedule_date,
end_date=datetime.now(tz=timezone.utc),
)
)
span.record_exception(reschedule)
span.set_status(
Status(StatusCode.ERROR, description=f"Exception: {type(reschedule).__name__}")
)
sys.exit(0)
with BundleVersionLock(
bundle_name=ti.bundle_instance.name,
bundle_version=ti.bundle_instance.version,
):
state, _, error = run(ti, context, log)
context["exception"] = error
# run() funnels every failure path into `error` rather than
# re-raising, so the worker span never sees the exception via
# propagation. Mark it here, where the worker span is in scope
# regardless of trace detail level.
if error is not None:
span.record_exception(error)
span.set_status(
Status(StatusCode.ERROR, description=f"Exception: {type(error).__name__}")
)
finalize(ti, state, context, log, error)
# If run() couldn't deliver a FAILED / UP_FOR_RETRY terminal
# state to the supervisor, fail closed now — finalize() has
# already run, so callbacks and listeners observed the state.
if getattr(ti, "_terminal_state_send_failed", False):
sys.exit(1)
except KeyboardInterrupt as e:
log.exception("Ctrl-c hit")
span.record_exception(e)
span.set_status(Status(StatusCode.ERROR, description=f"Exception: {type(e).__name__}"))
sys.exit(2)
except Exception as e:
log.exception("Top level error")
span.record_exception(e)
span.set_status(Status(StatusCode.ERROR, description=f"Exception: {type(e).__name__}"))
sys.exit(1)
finally:
# Ensure the request socket is closed on the child side in all circumstances
# before the process fully terminates.
if SUPERVISOR_COMMS and SUPERVISOR_COMMS.socket:
with suppress(Exception):
SUPERVISOR_COMMS.socket.close()
def reinit_supervisor_comms() -> None:
"""
Re-initialize supervisor comms and logging channel in subprocess.
This is not needed for most cases, but is used when either we re-launch the process via sudo for
run_as_user, or from inside the python code in a virtualenv (et al.) operator to re-connect so those tasks
can continue to access variables etc.
"""
import socket
if "SUPERVISOR_COMMS" not in globals():
global SUPERVISOR_COMMS
fd = int(os.environ.get("__AIRFLOW_SUPERVISOR_FD", "0"))
SUPERVISOR_COMMS = CommsDecoder[ToTask, ToSupervisor](log=log, socket=socket.socket(fileno=fd))
logs = SUPERVISOR_COMMS.send(ResendLoggingFD())
if isinstance(logs, SentFDs):
from airflow.sdk.log import configure_logging
log_io = os.fdopen(logs.fds[0], "wb", buffering=0)
configure_logging(json_output=True, output=log_io, sending_to_supervisor=True)
else:
print("Unable to re-configure logging after sudo, we didn't get an FD", file=sys.stderr)
if __name__ == "__main__":
main()
Reference in new issue View Git Blame Copy Permalink