From 3a182b4b70a3d43d8e0eae7076b6f2024a9e3eca Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Fri, 4 Nov 2022 05:11:10 +0800 Subject: [PATCH 1/3] Support mapped task group expansion at runtime This is relatively simple in concept, mostly just involving moving things around and minor refactoring so ti expansion now works for both MappedOperator (already does) and BaseOperator in a MappedTaskGroup. --- airflow/models/abstractoperator.py | 117 +++++++++++++++++++++++++++++ airflow/models/dagrun.py | 93 ++++++++++++++--------- airflow/models/mappedoperator.py | 117 ----------------------------- tests/models/test_dagrun.py | 38 ++++++++++ 4 files changed, 213 insertions(+), 152 deletions(-) diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py index efdbedcfdf444..98dd23548243d 100644 --- a/airflow/models/abstractoperator.py +++ b/airflow/models/abstractoperator.py @@ -26,12 +26,14 @@ from airflow.compat.functools import cache, cached_property from airflow.configuration import conf from airflow.exceptions import AirflowException +from airflow.models.expandinput import NotFullyPopulated from airflow.models.taskmixin import DAGNode from airflow.utils.context import Context from airflow.utils.helpers import render_template_as_native, render_template_to_string from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.mixins import ResolveMixin from airflow.utils.session import NEW_SESSION, provide_session +from airflow.utils.state import State, TaskInstanceState from airflow.utils.task_group import MappedTaskGroup from airflow.utils.trigger_rule import TriggerRule from airflow.utils.weight_rule import WeightRule @@ -422,6 +424,121 @@ def get_mapped_ti_count(self, run_id: str, *, session: Session) -> int: counts = (g.get_mapped_ti_count(run_id, session=session) for g in mapped_task_groups) return functools.reduce(operator.mul, counts) + def expand_mapped_task(self, run_id: str, *, session: Session) -> tuple[Sequence[TaskInstance], int]: + """Create the mapped task instances for mapped task. + + :raise NotMapped: If this task does not need expansion. + :return: The newly created mapped task instances (if any) in ascending + order by map index, and the maximum map index value. + """ + from sqlalchemy import func, or_ + + from airflow.models.baseoperator import BaseOperator + from airflow.models.mappedoperator import MappedOperator + from airflow.models.taskinstance import TaskInstance + from airflow.settings import task_instance_mutation_hook + + if not isinstance(self, (BaseOperator, MappedOperator)): + raise RuntimeError(f"cannot expand unrecognized operator type {type(self).__name__}") + + try: + total_length: int | None = self.get_mapped_ti_count(run_id, session=session) + except NotFullyPopulated as e: + # It's possible that the upstream tasks are not yet done, but we + # don't have upstream of upstreams in partial DAGs (possible in the + # mini-scheduler), so we ignore this exception. + if not self.dag or not self.dag.partial: + self.log.error( + "Cannot expand %r for run %s; missing upstream values: %s", + self, + run_id, + sorted(e.missing), + ) + total_length = None + + state: TaskInstanceState | None = None + unmapped_ti: TaskInstance | None = ( + session.query(TaskInstance) + .filter( + TaskInstance.dag_id == self.dag_id, + TaskInstance.task_id == self.task_id, + TaskInstance.run_id == run_id, + TaskInstance.map_index == -1, + or_(TaskInstance.state.in_(State.unfinished), TaskInstance.state.is_(None)), + ) + .one_or_none() + ) + + all_expanded_tis: list[TaskInstance] = [] + + if unmapped_ti: + # The unmapped task instance still exists and is unfinished, i.e. we + # haven't tried to run it before. + if total_length is None: + # If the DAG is partial, it's likely that the upstream tasks + # are not done yet, so the task can't fail yet. + if not self.dag or not self.dag.partial: + unmapped_ti.state = TaskInstanceState.UPSTREAM_FAILED + indexes_to_map: Iterable[int] = () + elif total_length < 1: + # If the upstream maps this to a zero-length value, simply mark + # the unmapped task instance as SKIPPED (if needed). + self.log.info( + "Marking %s as SKIPPED since the map has %d values to expand", + unmapped_ti, + total_length, + ) + unmapped_ti.state = TaskInstanceState.SKIPPED + indexes_to_map = () + else: + # Otherwise convert this into the first mapped index, and create + # TaskInstance for other indexes. + unmapped_ti.map_index = 0 + self.log.debug("Updated in place to become %s", unmapped_ti) + all_expanded_tis.append(unmapped_ti) + indexes_to_map = range(1, total_length) + state = unmapped_ti.state + elif not total_length: + # Nothing to fixup. + indexes_to_map = () + else: + # Only create "missing" ones. + current_max_mapping = ( + session.query(func.max(TaskInstance.map_index)) + .filter( + TaskInstance.dag_id == self.dag_id, + TaskInstance.task_id == self.task_id, + TaskInstance.run_id == run_id, + ) + .scalar() + ) + indexes_to_map = range(current_max_mapping + 1, total_length) + + for index in indexes_to_map: + # TODO: Make more efficient with bulk_insert_mappings/bulk_save_mappings. + ti = TaskInstance(self, run_id=run_id, map_index=index, state=state) + self.log.debug("Expanding TIs upserted %s", ti) + task_instance_mutation_hook(ti) + ti = session.merge(ti) + ti.refresh_from_task(self) # session.merge() loses task information. + all_expanded_tis.append(ti) + + # Coerce the None case to 0 -- these two are almost treated identically, + # except the unmapped ti (if exists) is marked to different states. + total_expanded_ti_count = total_length or 0 + + # Set to "REMOVED" any (old) TaskInstances with map indices greater + # than the current map value + session.query(TaskInstance).filter( + TaskInstance.dag_id == self.dag_id, + TaskInstance.task_id == self.task_id, + TaskInstance.run_id == run_id, + TaskInstance.map_index >= total_expanded_ti_count, + ).update({TaskInstance.state: TaskInstanceState.REMOVED}) + + session.flush() + return all_expanded_tis, total_expanded_ti_count - 1 + def render_template_fields( self, context: Context, diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index 07e1f1f489fee..fdf5f2277e0fe 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -51,7 +51,6 @@ from airflow.models.abstractoperator import NotMapped from airflow.models.base import Base, StringID from airflow.models.expandinput import NotFullyPopulated -from airflow.models.mappedoperator import MappedOperator from airflow.models.taskinstance import TaskInstance as TI from airflow.models.tasklog import LogTemplate from airflow.stats import Stats @@ -730,6 +729,27 @@ def _get_ready_tis( finished_tis=finished_tis, ) + def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None: + """Try to expand the ti, if needed. + + If the ti needs expansion, newly created task instances are + returned. The original ti is modified in-place and assigned the + ``map_index`` of 0. + + If the ti does not need expansion, either because the task is not + mapped, or has already been expanded, *None* is returned. + """ + if ti.map_index >= 0: # Already expanded, we're good. + return None + try: + expanded_tis, _ = ti.task.expand_mapped_task(self.run_id, session=session) + except NotMapped: # Not a mapped task, nothing needed. + return None + if expanded_tis: + assert expanded_tis[0] is ti + return expanded_tis[1:] + return () + # Check dependencies. expansion_happened = False for schedulable in itertools.chain(schedulable_tis, additional_tis): @@ -737,23 +757,19 @@ def _get_ready_tis( if not schedulable.are_dependencies_met(session=session, dep_context=dep_context): old_states[schedulable.key] = old_state continue - # If schedulable is from a mapped task, but not yet expanded, do it - # now. This is called in two places: First and ideally in the mini - # scheduler at the end of LocalTaskJob, and then as an "expansion of - # last resort" in the scheduler to ensure that the mapped task is - # correctly expanded before executed. - if schedulable.map_index < 0 and isinstance(schedulable.task, MappedOperator): - expanded_tis, _ = schedulable.task.expand_mapped_task(self.run_id, session=session) - if expanded_tis: - assert expanded_tis[0] is schedulable - additional_tis.extend(expanded_tis[1:]) - expansion_happened = True + # If schedulable is not yet expanded, try doing it now. This is + # called in two places: First and ideally in the mini scheduler at + # the end of LocalTaskJob, and then as an "expansion of last resort" + # in the scheduler to ensure that the mapped task is correctly + # expanded before executed. Also see _revise_map_indexes_if_mapped + # docstring for additional information. + if schedulable.map_index < 0: + new_tis = _expand_mapped_task_if_needed(schedulable) + if new_tis is not None: + additional_tis.extend(new_tis) + expansion_happened = True if schedulable.state in SCHEDULEABLE_STATES: - task = schedulable.task - if isinstance(task, MappedOperator): - # Ensure the task indexes are complete - created = self._revise_mapped_task_indexes(task, session=session) - ready_tis.extend(created) + ready_tis.extend(self._revise_map_indexes_if_mapped(schedulable.task, session=session)) ready_tis.append(schedulable) # Check if any ti changed state @@ -1080,14 +1096,22 @@ def _create_task_instances( # TODO[HA]: We probably need to savepoint this so we can keep the transaction alive. session.rollback() - def _revise_mapped_task_indexes(self, task: MappedOperator, session: Session) -> Iterable[TI]: - """Check if task increased or reduced in length and handle appropriately""" + def _revise_map_indexes_if_mapped(self, task: Operator, *, session: Session) -> Iterator[TI]: + """Check if task increased or reduced in length and handle appropriately. + + Currently missing tis are created and returned if possible. Expansion + only happens if depended upstreams are all ready; if not all of them + are, we delay expansion to the "last resort". See comments at the call + site for more details. + """ from airflow.settings import task_instance_mutation_hook try: total_length = task.get_mapped_ti_count(self.run_id, session=session) - except NotFullyPopulated: # Upstreams not ready, don't need to revise this yet. - return [] + except NotMapped: + return # Not a mapped task, don't need to do anything. + except NotFullyPopulated: + return # Upstreams not ready, don't need to revise this yet. query = session.query(TI.map_index).filter( TI.dag_id == self.dag_id, @@ -1095,20 +1119,9 @@ def _revise_mapped_task_indexes(self, task: MappedOperator, session: Session) -> TI.run_id == self.run_id, ) existing_indexes = {i for (i,) in query} - missing_indexes = set(range(total_length)).difference(existing_indexes) + removed_indexes = existing_indexes.difference(range(total_length)) - created_tis = [] - - if missing_indexes: - for index in missing_indexes: - ti = TI(task, run_id=self.run_id, map_index=index, state=None) - self.log.debug("Expanding TIs upserted %s", ti) - task_instance_mutation_hook(ti) - ti = session.merge(ti) - ti.refresh_from_task(task) - session.flush() - created_tis.append(ti) - elif removed_indexes: + if removed_indexes: session.query(TI).filter( TI.dag_id == self.dag_id, TI.task_id == task.task_id, @@ -1116,7 +1129,17 @@ def _revise_mapped_task_indexes(self, task: MappedOperator, session: Session) -> TI.map_index.in_(removed_indexes), ).update({TI.state: TaskInstanceState.REMOVED}) session.flush() - return created_tis + + for index in range(total_length): + if index in existing_indexes: + continue + ti = TI(task, run_id=self.run_id, map_index=index, state=None) + self.log.debug("Expanding TIs upserted %s", ti) + task_instance_mutation_hook(ti) + ti = session.merge(ti) + ti.refresh_from_task(task) + session.flush() + yield ti @staticmethod def get_run(session: Session, dag_id: str, execution_date: datetime) -> DagRun | None: diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py index 64de290351c8f..174d40074b4c2 100644 --- a/airflow/models/mappedoperator.py +++ b/airflow/models/mappedoperator.py @@ -27,7 +27,6 @@ import attr import pendulum -from sqlalchemy import func, or_ from sqlalchemy.orm.session import Session from airflow import settings @@ -51,7 +50,6 @@ DictOfListsExpandInput, ExpandInput, ListOfDictsExpandInput, - NotFullyPopulated, OperatorExpandArgument, OperatorExpandKwargsArgument, is_mappable, @@ -65,7 +63,6 @@ from airflow.utils.context import Context, context_update_for_unmapped from airflow.utils.helpers import is_container, prevent_duplicates from airflow.utils.operator_resources import Resources -from airflow.utils.state import State, TaskInstanceState from airflow.utils.trigger_rule import TriggerRule from airflow.utils.types import NOTSET @@ -75,7 +72,6 @@ from airflow.models.baseoperator import BaseOperator, BaseOperatorLink from airflow.models.dag import DAG from airflow.models.operator import Operator - from airflow.models.taskinstance import TaskInstance from airflow.models.xcom_arg import XComArg from airflow.utils.task_group import TaskGroup @@ -608,119 +604,6 @@ def _get_specified_expand_input(self) -> ExpandInput: """Input received from the expand call on the operator.""" return getattr(self, self._expand_input_attr) - def expand_mapped_task(self, run_id: str, *, session: Session) -> tuple[Sequence[TaskInstance], int]: - """Create the mapped task instances for mapped task. - - :return: The newly created mapped TaskInstances (if any) in ascending order by map index, and the - maximum map_index. - """ - from airflow.models.taskinstance import TaskInstance - from airflow.settings import task_instance_mutation_hook - - total_length: int | None - try: - total_length = self._get_specified_expand_input().get_total_map_length(run_id, session=session) - except NotFullyPopulated as e: - total_length = None - # partial dags comes from the mini scheduler. It's - # possible that the upstream tasks are not yet done, - # but we don't have upstream of upstreams in partial dags, - # so we ignore this exception. - if not self.dag or not self.dag.partial: - self.log.error( - "Cannot expand %r for run %s; missing upstream values: %s", - self, - run_id, - sorted(e.missing), - ) - - state: TaskInstanceState | None = None - unmapped_ti: TaskInstance | None = ( - session.query(TaskInstance) - .filter( - TaskInstance.dag_id == self.dag_id, - TaskInstance.task_id == self.task_id, - TaskInstance.run_id == run_id, - TaskInstance.map_index == -1, - or_(TaskInstance.state.in_(State.unfinished), TaskInstance.state.is_(None)), - ) - .one_or_none() - ) - - all_expanded_tis: list[TaskInstance] = [] - - if unmapped_ti: - # The unmapped task instance still exists and is unfinished, i.e. we - # haven't tried to run it before. - if total_length is None: - if self.dag and self.dag.partial: - # If the DAG is partial, it's likely that the upstream tasks - # are not done yet, so we do nothing - indexes_to_map: Iterable[int] = () - else: - # If the map length cannot be calculated (due to unavailable - # upstream sources), fail the unmapped task. - unmapped_ti.state = TaskInstanceState.UPSTREAM_FAILED - indexes_to_map = () - elif total_length < 1: - # If the upstream maps this to a zero-length value, simply mark - # the unmapped task instance as SKIPPED (if needed). - self.log.info( - "Marking %s as SKIPPED since the map has %d values to expand", - unmapped_ti, - total_length, - ) - unmapped_ti.state = TaskInstanceState.SKIPPED - indexes_to_map = () - else: - # Otherwise convert this into the first mapped index, and create - # TaskInstance for other indexes. - unmapped_ti.map_index = 0 - self.log.debug("Updated in place to become %s", unmapped_ti) - all_expanded_tis.append(unmapped_ti) - indexes_to_map = range(1, total_length) - state = unmapped_ti.state - elif not total_length: - # Nothing to fixup. - indexes_to_map = () - else: - # Only create "missing" ones. - current_max_mapping = ( - session.query(func.max(TaskInstance.map_index)) - .filter( - TaskInstance.dag_id == self.dag_id, - TaskInstance.task_id == self.task_id, - TaskInstance.run_id == run_id, - ) - .scalar() - ) - indexes_to_map = range(current_max_mapping + 1, total_length) - - for index in indexes_to_map: - # TODO: Make more efficient with bulk_insert_mappings/bulk_save_mappings. - ti = TaskInstance(self, run_id=run_id, map_index=index, state=state) - self.log.debug("Expanding TIs upserted %s", ti) - task_instance_mutation_hook(ti) - ti = session.merge(ti) - ti.refresh_from_task(self) # session.merge() loses task information. - all_expanded_tis.append(ti) - - # Coerce the None case to 0 -- these two are almost treated identically, - # except the unmapped ti (if exists) is marked to different states. - total_expanded_ti_count = total_length or 0 - - # Set to "REMOVED" any (old) TaskInstances with map indices greater - # than the current map value - session.query(TaskInstance).filter( - TaskInstance.dag_id == self.dag_id, - TaskInstance.task_id == self.task_id, - TaskInstance.run_id == run_id, - TaskInstance.map_index >= total_expanded_ti_count, - ).update({TaskInstance.state: TaskInstanceState.REMOVED}) - - session.flush() - return all_expanded_tis, total_expanded_ti_count - 1 - def prepare_for_execution(self) -> MappedOperator: # Since a mapped operator cannot be used for execution, and an unmapped # BaseOperator needs to be created later (see render_template_fields), diff --git a/tests/models/test_dagrun.py b/tests/models/test_dagrun.py index 6295dbcf90c16..f88d4c1db9de6 100644 --- a/tests/models/test_dagrun.py +++ b/tests/models/test_dagrun.py @@ -2009,3 +2009,41 @@ def execute(self, context): ti.run() assert sorted(results) == expected + + +def test_mapped_task_group_expands(dag_maker, session): + with dag_maker(session=session): + + @task_group + def tg(x, y): + return MockOperator(task_id="task_2", arg1=x, arg2=y) + + task_1 = BaseOperator(task_id="task_1") + tg.expand(x=task_1.output, y=[1, 2, 3]) + + dr: DagRun = dag_maker.create_dagrun() + + # Not expanding task_2 yet since it depends on result from task_1. + decision = dr.task_instance_scheduling_decisions(session=session) + assert {(ti.task_id, ti.map_index, ti.state) for ti in decision.tis} == { + ("task_1", -1, None), + ("tg.task_2", -1, None), + } + + # Simulate task_1 execution to produce TaskMap. + (ti_1,) = decision.schedulable_tis + assert ti_1.task_id == "task_1" + ti_1.state = TaskInstanceState.SUCCESS + session.add(TaskMap.from_task_instance_xcom(ti_1, ["a", "b"])) + session.flush() + + # Now task_2 in mapped tagk group is expanded. + decision = dr.task_instance_scheduling_decisions(session=session) + assert {(ti.task_id, ti.map_index, ti.state) for ti in decision.schedulable_tis} == { + ("tg.task_2", 0, None), + ("tg.task_2", 1, None), + ("tg.task_2", 2, None), + ("tg.task_2", 3, None), + ("tg.task_2", 4, None), + ("tg.task_2", 5, None), + } From 2d6dfdc3e7edfca4307c8852f9637c7752fd87ef Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 7 Nov 2022 14:27:39 +0800 Subject: [PATCH 2/3] Correctly set up XComArg deps in MappedArgument --- airflow/models/expandinput.py | 21 +++++++++++++++++++++ airflow/models/mappedoperator.py | 5 ++--- airflow/models/param.py | 21 +++++++++++++-------- airflow/models/xcom_arg.py | 30 ++++++++++++++---------------- airflow/utils/mixins.py | 16 ++++++++++++++++ 5 files changed, 66 insertions(+), 27 deletions(-) diff --git a/airflow/models/expandinput.py b/airflow/models/expandinput.py index b817da38dfcde..a7ab9875a02d3 100644 --- a/airflow/models/expandinput.py +++ b/airflow/models/expandinput.py @@ -32,6 +32,7 @@ if TYPE_CHECKING: from sqlalchemy.orm import Session + from airflow.models.operator import Operator from airflow.models.xcom_arg import XComArg ExpandInput = Union["DictOfListsExpandInput", "ListOfDictsExpandInput"] @@ -61,6 +62,9 @@ def get_task_map_length(self, run_id: str, *, session: Session) -> int | None: # This simply marks the value as un-expandable at parse-time. return None + def iter_references(self) -> Iterable[tuple[Operator, str]]: + yield from self._input.iter_references() + @provide_session def resolve(self, context: Context, *, session: Session = NEW_SESSION) -> Any: data, _ = self._input.resolve(context, session=session) @@ -185,6 +189,13 @@ def _find_index_for_this_field(index: int) -> int: return k, v raise IndexError(f"index {map_index} is over mapped length") + def iter_references(self) -> Iterable[tuple[Operator, str]]: + from airflow.models.xcom_arg import XComArg + + for x in self.value.values(): + if isinstance(x, XComArg): + yield from x.iter_references() + def resolve(self, context: Context, session: Session) -> tuple[Mapping[str, Any], set[int]]: data = {k: self._expand_mapped_field(k, v, context, session=session) for k, v in self.value.items()} literal_keys = {k for k, _ in self._iter_parse_time_resolved_kwargs()} @@ -219,6 +230,16 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int: raise NotFullyPopulated({"expand_kwargs() argument"}) return length + def iter_references(self) -> Iterable[tuple[Operator, str]]: + from airflow.models.xcom_arg import XComArg + + if isinstance(self.value, XComArg): + yield from self.value.iter_references() + else: + for x in self.value: + if isinstance(x, XComArg): + yield from x.iter_references() + def resolve(self, context: Context, session: Session) -> tuple[Mapping[str, Any], set[int]]: map_index = context["ti"].map_index if map_index < 0: diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py index 174d40074b4c2..aea0ce084b253 100644 --- a/airflow/models/mappedoperator.py +++ b/airflow/models/mappedoperator.py @@ -614,9 +614,8 @@ def iter_mapped_dependencies(self) -> Iterator[Operator]: """Upstream dependencies that provide XComs used by this task for task mapping.""" from airflow.models.xcom_arg import XComArg - for ref in XComArg.iter_xcom_args(self._get_specified_expand_input()): - for operator, _ in ref.iter_references(): - yield operator + for operator, _ in XComArg.iter_xcom_references(self._get_specified_expand_input()): + yield operator @cache def get_parse_time_mapped_ti_count(self) -> int: diff --git a/airflow/models/param.py b/airflow/models/param.py index be625227c531b..4a087e2a12e56 100644 --- a/airflow/models/param.py +++ b/airflow/models/param.py @@ -21,7 +21,7 @@ import json import logging import warnings -from typing import TYPE_CHECKING, Any, ItemsView, MutableMapping, ValuesView +from typing import TYPE_CHECKING, Any, ItemsView, Iterable, MutableMapping, ValuesView from airflow.exceptions import AirflowException, ParamValidationError, RemovedInAirflow3Warning from airflow.utils.context import Context @@ -233,15 +233,17 @@ def validate(self) -> dict[str, Any]: class DagParam(ResolveMixin): - """ - Class that represents a DAG run parameter & binds a simple Param object to a name within a DAG instance, - so that it can be resolved during the run time via ``{{ context }}`` dictionary. The ideal use case of - this class is to implicitly convert args passed to a method which is being decorated by ``@dag`` keyword. + """DAG run parameter reference. + + This binds a simple Param object to a name within a DAG instance, so that it + can be resolved during the runtime via the ``{{ context }}`` dictionary. The + ideal use case of this class is to implicitly convert args passed to a + method decorated by ``@dag``. - It can be used to parameterize your dags. You can overwrite its value by setting it on conf - when you trigger your DagRun. + It can be used to parameterize a DAG. You can overwrite its value by setting + it on conf when you trigger your DagRun. - This can also be used in templates by accessing ``{{context.params}}`` dictionary. + This can also be used in templates by accessing ``{{ context.params }}``. **Example**: @@ -259,6 +261,9 @@ def __init__(self, current_dag: DAG, name: str, default: Any = NOTSET): self._name = name self._default = default + def iter_references(self) -> Iterable[tuple[Operator, str]]: + return () + def resolve(self, context: Context) -> Any: """Pull DagParam value from DagRun context. This method is run during ``op.execute()``.""" with contextlib.suppress(KeyError): diff --git a/airflow/models/xcom_arg.py b/airflow/models/xcom_arg.py index 2a95c32863ad6..2fb43af34bb4d 100644 --- a/airflow/models/xcom_arg.py +++ b/airflow/models/xcom_arg.py @@ -91,23 +91,23 @@ def __new__(cls, *args, **kwargs) -> XComArg: return super().__new__(cls) @staticmethod - def iter_xcom_args(arg: Any) -> Iterator[XComArg]: - """Return XComArg instances in an arbitrary value. + def iter_xcom_references(arg: Any) -> Iterator[tuple[Operator, str]]: + """Return XCom references in an arbitrary value. Recursively traverse ``arg`` and look for XComArg instances in any collection objects, and instances with ``template_fields`` set. """ - if isinstance(arg, XComArg): - yield arg + if isinstance(arg, ResolveMixin): + yield from arg.iter_references() elif isinstance(arg, (tuple, set, list)): for elem in arg: - yield from XComArg.iter_xcom_args(elem) + yield from XComArg.iter_xcom_references(elem) elif isinstance(arg, dict): for elem in arg.values(): - yield from XComArg.iter_xcom_args(elem) + yield from XComArg.iter_xcom_references(elem) elif isinstance(arg, AbstractOperator): - for elem in arg.template_fields: - yield from XComArg.iter_xcom_args(elem) + for attr in arg.template_fields: + yield from XComArg.iter_xcom_references(getattr(arg, attr)) @staticmethod def apply_upstream_relationship(op: Operator, arg: Any): @@ -117,9 +117,8 @@ def apply_upstream_relationship(op: Operator, arg: Any): collections objects and classes decorated with ``template_fields``), and sets the relationship to ``op`` on any found. """ - for ref in XComArg.iter_xcom_args(arg): - for operator, _ in ref.iter_references(): - op.set_upstream(operator) + for operator, _ in XComArg.iter_xcom_references(arg): + op.set_upstream(operator) @property def roots(self) -> list[DAGNode]: @@ -173,10 +172,6 @@ def _deserialize(cls, data: dict[str, Any], dag: DAG) -> XComArg: """ raise NotImplementedError() - def iter_references(self) -> Iterator[tuple[Operator, str]]: - """Iterate through (operator, key) references.""" - raise NotImplementedError() - def map(self, f: Callable[[Any], Any]) -> MapXComArg: return MapXComArg(self, [f]) @@ -196,7 +191,10 @@ def get_task_map_length(self, run_id: str, *, session: Session) -> int | None: def resolve(self, context: Context, session: Session = NEW_SESSION) -> Any: """Pull XCom value. - This should only be called during ``op.execute()`` in respectable context. + This should only be called during ``op.execute()`` in respectable + context. Note that although the ``ResolveMixin`` parent mixin also has a + ``resolve`` protocol, this adds the optional ``session`` argument that + some XComArg usages need. :meta private: """ diff --git a/airflow/utils/mixins.py b/airflow/utils/mixins.py index 12ae526a09afe..7cef6cf6b226c 100644 --- a/airflow/utils/mixins.py +++ b/airflow/utils/mixins.py @@ -24,6 +24,9 @@ from airflow.configuration import conf from airflow.utils.context import Context +if typing.TYPE_CHECKING: + from airflow.models.operator import Operator + class MultiprocessingStartMethodMixin: """Convenience class to add support for different types of multiprocessing.""" @@ -45,5 +48,18 @@ def _get_multiprocessing_start_method(self) -> str: class ResolveMixin: """A runtime-resolved value.""" + def iter_references(self) -> typing.Iterable[tuple[Operator, str]]: + """Find underlying XCom references this contains. + + This is used by the DAG parser to recursively find task dependencies. + + :meta private: + """ + raise NotImplementedError + def resolve(self, context: Context) -> typing.Any: + """Resolve this value for runtime. + + :meta private: + """ raise NotImplementedError From a17e4103113903e0f8628ea4b2fe6417928c18a0 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 17 Nov 2022 15:27:51 +0800 Subject: [PATCH 3/3] Improve docstrings and comments for code readers --- airflow/models/abstractoperator.py | 4 ++-- airflow/models/dagrun.py | 8 ++++---- airflow/models/xcom_arg.py | 9 +++++---- 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py index 98dd23548243d..9ea1873b98ea8 100644 --- a/airflow/models/abstractoperator.py +++ b/airflow/models/abstractoperator.py @@ -527,8 +527,8 @@ def expand_mapped_task(self, run_id: str, *, session: Session) -> tuple[Sequence # except the unmapped ti (if exists) is marked to different states. total_expanded_ti_count = total_length or 0 - # Set to "REMOVED" any (old) TaskInstances with map indices greater - # than the current map value + # Any (old) task instances with inapplicable indexes (>= the total + # number we need) are set to "REMOVED". session.query(TaskInstance).filter( TaskInstance.dag_id == self.dag_id, TaskInstance.task_id == self.task_id, diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index fdf5f2277e0fe..d4f5a3be42de3 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -1099,10 +1099,10 @@ def _create_task_instances( def _revise_map_indexes_if_mapped(self, task: Operator, *, session: Session) -> Iterator[TI]: """Check if task increased or reduced in length and handle appropriately. - Currently missing tis are created and returned if possible. Expansion - only happens if depended upstreams are all ready; if not all of them - are, we delay expansion to the "last resort". See comments at the call - site for more details. + Task instances that do not already exist are created and returned if + possible. Expansion only happens if all upstreams are ready; otherwise + we delay expansion to the "last resort". See comments at the call site + for more details. """ from airflow.settings import task_instance_mutation_hook diff --git a/airflow/models/xcom_arg.py b/airflow/models/xcom_arg.py index 2fb43af34bb4d..d43c5dffe040b 100644 --- a/airflow/models/xcom_arg.py +++ b/airflow/models/xcom_arg.py @@ -191,10 +191,11 @@ def get_task_map_length(self, run_id: str, *, session: Session) -> int | None: def resolve(self, context: Context, session: Session = NEW_SESSION) -> Any: """Pull XCom value. - This should only be called during ``op.execute()`` in respectable - context. Note that although the ``ResolveMixin`` parent mixin also has a - ``resolve`` protocol, this adds the optional ``session`` argument that - some XComArg usages need. + This should only be called during ``op.execute()`` with an appropriate + context (e.g. generated from ``TaskInstance.get_template_context()``). + Although the ``ResolveMixin`` parent mixin also has a ``resolve`` + protocol, this adds the optional ``session`` argument that some of the + subclasses need. :meta private: """