From ab13a0d4cc46166914e7afcd82c5413baa3827ca Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Thu, 20 Oct 2022 16:35:48 +0800 Subject: [PATCH 1/2] Expand tasks in mapped group at parse time This involves a few things: * The mapped task group needs to implement parse-time and run-time ti count accessors, similar to the mapped operator. * Since a non-mapped operator can be in a mapped task group and thus needs to be expanded, the parse-time and run-time ti count accessors are "raised" from MappedOperator to AbstractOperator (so they can be used by BaseOperator). These now also considers an operator's parent task groups. * If a ti count accessor is called on a non-mapped operator that's not in any mapped task group, the accessor should return a special value of a sort. Since None is already used, we raise an exception in this case instead (simply called NotMapped). For consistency, the previous case where None was returned (upstreams not available) is also changed to raise an exception instead (reusing NotFullyPopulated). This means the ti count accessors now always return an int, and raise an exception whenever the count is not available. * All of the above result in a coherant interface in AbstractOperator that can be used in DagRun to perform ti count revision. I also removed caching from get_mapped_ti_count since (judging from the various cache-clearing calls) it seems to be cause more trouble than benefits. --- airflow/models/abstractoperator.py | 54 ++++++++++++++++++- airflow/models/dagrun.py | 84 ++++++++++++++---------------- airflow/models/expandinput.py | 9 ++-- airflow/models/mappedoperator.py | 41 ++++++--------- airflow/utils/task_group.py | 31 +++++++++++ tests/models/test_taskinstance.py | 7 +-- tests/test_utils/mapping.py | 1 - 7 files changed, 147 insertions(+), 80 deletions(-) diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py index a7e3380dacc2c..efdbedcfdf444 100644 --- a/airflow/models/abstractoperator.py +++ b/airflow/models/abstractoperator.py @@ -18,10 +18,12 @@ from __future__ import annotations import datetime +import functools import inspect +import operator from typing import TYPE_CHECKING, Any, Callable, ClassVar, Collection, Iterable, Iterator, Sequence -from airflow.compat.functools import cached_property +from airflow.compat.functools import cache, cached_property from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.models.taskmixin import DAGNode @@ -30,6 +32,7 @@ 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.task_group import MappedTaskGroup from airflow.utils.trigger_rule import TriggerRule from airflow.utils.weight_rule import WeightRule @@ -65,6 +68,10 @@ ) +class NotMapped(Exception): + """Raise if a task is neither mapped nor has any parent mapped groups.""" + + class AbstractOperator(LoggingMixin, DAGNode): """Common implementation for operators, including unmapped and mapped. @@ -276,6 +283,14 @@ def iter_mapped_dependants(self) -> Iterator[MappedOperator]: if any(p.node_id == self.node_id for p in downstream.iter_mapped_dependencies()) ) + def iter_mapped_task_groups(self) -> Iterator[MappedTaskGroup]: + """Return mapped task groups this task belongs to.""" + parent = self.task_group + while parent is not None: + if isinstance(parent, MappedTaskGroup): + yield parent + parent = parent.task_group + def unmap(self, resolve: None | dict[str, Any] | tuple[Context, Session]) -> BaseOperator: """Get the "normal" operator from current abstract operator. @@ -370,6 +385,43 @@ def get_extra_links(self, ti: TaskInstance, link_name: str) -> str | None: return link.get_link(self.unmap(None), ti.dag_run.logical_date) # type: ignore[misc] return link.get_link(self.unmap(None), ti_key=ti.key) + @cache + def get_parse_time_mapped_ti_count(self) -> int: + """Number of mapped task instances that can be created on DAG run creation. + + This only considers literal mapped arguments, and would return *None* + when any non-literal values are used for mapping. + + :raise NotFullyPopulated: If non-literal mapped arguments are encountered. + :raise NotMapped: If the operator is neither mapped, nor has any parent + mapped task groups. + :return: Total number of mapped TIs this task should have. + """ + mapped_task_groups = list(self.iter_mapped_task_groups()) + if not mapped_task_groups: + raise NotMapped + counts = (g.get_parse_time_mapped_ti_count() for g in mapped_task_groups) + return functools.reduce(operator.mul, counts) + + def get_mapped_ti_count(self, run_id: str, *, session: Session) -> int: + """Number of mapped TaskInstances that can be created at run time. + + This considers both literal and non-literal mapped arguments, and the + result is therefore available when all depended tasks have finished. The + return value should be identical to ``parse_time_mapped_ti_count`` if + all mapped arguments are literal. + + :raise NotFullyPopulated: If upstream tasks are not all complete yet. + :raise NotMapped: If the operator is neither mapped, nor has any parent + mapped task groups. + :return: Total number of mapped TIs this task should have. + """ + mapped_task_groups = list(self.iter_mapped_task_groups()) + if not mapped_task_groups: + raise NotMapped + counts = (g.get_mapped_ti_count(run_id, session=session) for g in mapped_task_groups) + return functools.reduce(operator.mul, counts) + def render_template_fields( self, context: Context, diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index 025f195b329bb..07e1f1f489fee 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -22,18 +22,7 @@ import warnings from collections import defaultdict from datetime import datetime -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Iterable, - Iterator, - NamedTuple, - Sequence, - TypeVar, - cast, - overload, -) +from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator, NamedTuple, Sequence, TypeVar, overload from sqlalchemy import ( Boolean, @@ -59,7 +48,9 @@ from airflow.callbacks.callback_requests import DagCallbackRequest from airflow.configuration import conf as airflow_conf from airflow.exceptions import AirflowException, RemovedInAirflow3Warning, TaskNotFound +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 @@ -907,7 +898,6 @@ def _check_for_removed_or_restored_tasks( for ti in tis: ti_mutation_hook(ti) task_ids.add(ti.task_id) - task = None try: task = dag.get_task(ti.task_id) @@ -925,28 +915,15 @@ def _check_for_removed_or_restored_tasks( ti.state = State.REMOVED continue - if not task.is_mapped: + try: + num_mapped_tis = task.get_parse_time_mapped_ti_count() + except NotMapped: continue - task = cast("MappedOperator", task) - num_mapped_tis = task.parse_time_mapped_ti_count - # Check if the number of mapped literals has changed and we need to mark this TI as removed - if num_mapped_tis is not None: - if ti.map_index >= num_mapped_tis: - self.log.debug( - "Removing task '%s' as the map_index is longer than the literal mapping list (%s)", - ti, - num_mapped_tis, - ) - ti.state = State.REMOVED - elif ti.map_index < 0: - self.log.debug("Removing the unmapped TI '%s' as the mapping can now be performed", ti) - ti.state = State.REMOVED - else: - # What if it is _now_ dynamically mapped, but wasn't before? - task.get_mapped_ti_count.cache_clear() # type: ignore[attr-defined] - total_length = task.get_mapped_ti_count(self.run_id, session=session) - - if total_length is None: + except NotFullyPopulated: + # What if it is _now_ dynamically mapped, but wasn't before? + try: + total_length = task.get_mapped_ti_count(self.run_id, session=session) + except NotFullyPopulated: # Not all upstreams finished, so we can't tell what should be here. Remove everything. if ti.map_index >= 0: self.log.debug( @@ -962,6 +939,18 @@ def _check_for_removed_or_restored_tasks( total_length, ) ti.state = State.REMOVED + else: + # Check if the number of mapped literals has changed and we need to mark this TI as removed. + if ti.map_index >= num_mapped_tis: + self.log.debug( + "Removing task '%s' as the map_index is longer than the literal mapping list (%s)", + ti, + num_mapped_tis, + ) + ti.state = State.REMOVED + elif ti.map_index < 0: + self.log.debug("Removing the unmapped TI '%s' as the mapping can now be performed", ti) + ti.state = State.REMOVED return task_ids @@ -1033,15 +1022,20 @@ def _create_tasks( :param tasks: Tasks to create jobs for in the DAG run :param task_creator: Function to create task instances """ + map_indexes: Iterable[int] for task in tasks: - if not task.is_mapped: - yield from task_creator(task, (-1,)) - continue - count = cast(MappedOperator, task).get_mapped_ti_count(self.run_id, session=session) - if count: - yield from task_creator(task, range(count)) - continue - yield from task_creator(task, (-1,)) + try: + count = task.get_mapped_ti_count(self.run_id, session=session) + except (NotMapped, NotFullyPopulated): + map_indexes = (-1,) + else: + if count: + map_indexes = range(count) + else: + # Make sure to always create at least one ti; this will be + # marked as REMOVED later at runtime. + map_indexes = (-1,) + yield from task_creator(task, map_indexes) def _create_task_instances( self, @@ -1090,9 +1084,9 @@ def _revise_mapped_task_indexes(self, task: MappedOperator, session: Session) -> """Check if task increased or reduced in length and handle appropriately""" from airflow.settings import task_instance_mutation_hook - task.get_mapped_ti_count.cache_clear() # type: ignore[attr-defined] - total_length = task.get_mapped_ti_count(self.run_id, session=session) - if total_length is None: # Upstreams not ready, don't need to revise this yet. + 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 [] query = session.query(TI.map_index).filter( diff --git a/airflow/models/expandinput.py b/airflow/models/expandinput.py index d8b13aa09e395..50094933bc597 100644 --- a/airflow/models/expandinput.py +++ b/airflow/models/expandinput.py @@ -78,12 +78,13 @@ def _iter_parse_time_resolved_kwargs(self) -> Iterable[tuple[str, Sized]]: return ((k, v) for k, v in self.value.items() if not isinstance(v, XComArg)) - def get_parse_time_mapped_ti_count(self) -> int | None: + def get_parse_time_mapped_ti_count(self) -> int: if not self.value: return 0 literal_values = [len(v) for _, v in self._iter_parse_time_resolved_kwargs()] if len(literal_values) != len(self.value): - return None # None-literal type encountered, so give up. + literal_keys = (k for k, _ in self._iter_parse_time_resolved_kwargs()) + raise NotFullyPopulated(set(self.value).difference(literal_keys)) return functools.reduce(operator.mul, literal_values, 1) def _get_map_lengths(self, run_id: str, *, session: Session) -> dict[str, int]: @@ -166,10 +167,10 @@ class ListOfDictsExpandInput(NamedTuple): value: OperatorExpandKwargsArgument - def get_parse_time_mapped_ti_count(self) -> int | None: + def get_parse_time_mapped_ti_count(self) -> int: if isinstance(self.value, collections.abc.Sized): return len(self.value) - return None + raise NotFullyPopulated({"expand_kwargs() argument"}) def get_total_map_length(self, run_id: str, *, session: Session) -> int: if isinstance(self.value, collections.abc.Sized): diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py index 2d5d00cf3dfe9..35f56279e76af 100644 --- a/airflow/models/mappedoperator.py +++ b/airflow/models/mappedoperator.py @@ -31,7 +31,7 @@ from sqlalchemy.orm.session import Session from airflow import settings -from airflow.compat.functools import cache, cached_property +from airflow.compat.functools import cache from airflow.exceptions import AirflowException, UnmappableOperator from airflow.models.abstractoperator import ( DEFAULT_IGNORE_FIRST_DEPENDS_ON_PAST, @@ -44,6 +44,7 @@ DEFAULT_TRIGGER_RULE, DEFAULT_WEIGHT_RULE, AbstractOperator, + NotMapped, TaskStateChangeCallback, ) from airflow.models.expandinput import ( @@ -724,34 +725,22 @@ def iter_mapped_dependencies(self) -> Iterator[Operator]: for operator, _ in ref.iter_references(): yield operator - @cached_property - def parse_time_mapped_ti_count(self) -> int | None: - """Number of mapped TaskInstances that can be created at DagRun create time. - - This only considers literal mapped arguments, and would return *None* - when any non-literal values are used for mapping. - - :return: None if non-literal mapped arg encountered, or the total - number of mapped TIs this task should have. - """ - return self._get_specified_expand_input().get_parse_time_mapped_ti_count() - @cache - def get_mapped_ti_count(self, run_id: str, *, session: Session) -> int | None: - """Number of mapped TaskInstances that can be created at run time. - - This considers both literal and non-literal mapped arguments, and the - result is therefore available when all depended tasks have finished. The - return value should be identical to ``parse_time_mapped_ti_count`` if - all mapped arguments are literal. + def get_parse_time_mapped_ti_count(self) -> int: + current_count = self._get_specified_expand_input().get_parse_time_mapped_ti_count() + try: + parent_count = super().get_parse_time_mapped_ti_count() + except NotMapped: + return current_count + return parent_count * current_count - :return: None if upstream tasks are not complete yet, or the total - number of mapped TIs this task should have. - """ + def get_mapped_ti_count(self, run_id: str, *, session: Session) -> int: + current_count = self._get_specified_expand_input().get_total_map_length(run_id, session=session) try: - return self._get_specified_expand_input().get_total_map_length(run_id, session=session) - except NotFullyPopulated: - return None + parent_count = super().get_mapped_ti_count(run_id, session=session) + except NotMapped: + return current_count + return parent_count * current_count def render_template_fields( self, diff --git a/airflow/utils/task_group.py b/airflow/utils/task_group.py index 0490b339915ff..1acf195890879 100644 --- a/airflow/utils/task_group.py +++ b/airflow/utils/task_group.py @@ -26,6 +26,7 @@ import weakref from typing import TYPE_CHECKING, Any, Generator, Sequence +from airflow.compat.functools import cache from airflow.exceptions import ( AirflowDagCycleException, AirflowException, @@ -37,6 +38,8 @@ from airflow.utils.helpers import validate_group_key if TYPE_CHECKING: + from sqlalchemy.orm import Session + from airflow.models.baseoperator import BaseOperator from airflow.models.dag import DAG from airflow.models.expandinput import ExpandInput @@ -469,6 +472,34 @@ def __init__(self, *, expand_input: ExpandInput, **kwargs: Any) -> None: super().__init__(**kwargs) self._expand_input = expand_input + @cache + def get_parse_time_mapped_ti_count(self) -> int: + """Number of instances a task in this group should be mapped to, when a DAG run is created. + + This only considers literal mapped arguments, and would return *None* + when any non-literal values are used for mapping. This also does not + account for nested mapped groups; only the number expanded due to this + specific group is returned, regardless of whether any of its parent + groups are mapped. + + :raise NotFullyPopulated: If any non-literal mapped arguments are encountered. + :return: The total number of mapped instances each task should have. + """ + return self._expand_input.get_parse_time_mapped_ti_count() + + def get_mapped_ti_count(self, run_id: str, *, session: Session) -> int: + """Number of instances a task in this group should be mapped to at run time. + + This considers both literal and non-literal mapped arguments, and the + result is therefore available when all depended tasks have finished. The + return value should be identical to ``parse_time_mapped_ti_count`` if + all mapped arguments are literal. + + :raise NotFullyPopulated: If upstream tasks are not all complete yet. + :return: Total number of mapped TIs this task should have. + """ + return self._expand_input.get_total_map_length(run_id, session=session) + class TaskGroupContext: """TaskGroup context is used to keep the current TaskGroup when TaskGroup is used as ContextManager.""" diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index d21203b2d752c..4130566a30e92 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -60,7 +60,7 @@ XCom, ) from airflow.models.dataset import DatasetDagRunQueue, DatasetEvent, DatasetModel -from airflow.models.expandinput import EXPAND_INPUT_EMPTY +from airflow.models.expandinput import EXPAND_INPUT_EMPTY, NotFullyPopulated from airflow.models.param import process_params from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskfail import TaskFail @@ -3376,7 +3376,8 @@ def show(a, b): ti.run() show_task = dag.get_task("show") - assert show_task.parse_time_mapped_ti_count is None + with pytest.raises(NotFullyPopulated): + assert show_task.get_parse_time_mapped_ti_count() mapped_tis, max_map_index = show_task.expand_mapped_task(dag_run.run_id, session=session) assert max_map_index + 1 == len(mapped_tis) == 4 @@ -3400,7 +3401,7 @@ def show(a, b): dag_run = dag_maker.create_dagrun() show_task = dag.get_task("show") - assert show_task.parse_time_mapped_ti_count == 6 + assert show_task.get_parse_time_mapped_ti_count() == 6 mapped_tis, max_map_index = show_task.expand_mapped_task(dag_run.run_id, session=session) assert len(mapped_tis) == 0 # Expanded at parse! assert max_map_index == 5 diff --git a/tests/test_utils/mapping.py b/tests/test_utils/mapping.py index 984446343cb89..15a57679cf024 100644 --- a/tests/test_utils/mapping.py +++ b/tests/test_utils/mapping.py @@ -42,4 +42,3 @@ def expand_mapped_task( session.flush() mapped.expand_mapped_task(run_id, session=session) - mapped.get_mapped_ti_count.cache_clear() # type: ignore[attr-defined] From b3b683780af10548bf4d126facb2cdfc3ed21c39 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Fri, 21 Oct 2022 15:34:39 +0800 Subject: [PATCH 2/2] Use type guards to detect run-time constructs --- airflow/decorators/base.py | 10 ++-- airflow/decorators/task_group.py | 21 ++------ airflow/models/expandinput.py | 75 ++++++++++++++++++++++------- airflow/models/mappedoperator.py | 4 +- airflow/typing_compat.py | 5 +- tests/decorators/test_task_group.py | 7 ++- tests/models/test_dagrun.py | 41 +++++++++++++++- 7 files changed, 110 insertions(+), 53 deletions(-) diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py index 715a5515ccaad..c27262f2b0725 100644 --- a/airflow/decorators/base.py +++ b/airflow/decorators/base.py @@ -55,13 +55,9 @@ ListOfDictsExpandInput, OperatorExpandArgument, OperatorExpandKwargsArgument, + is_mappable, ) -from airflow.models.mappedoperator import ( - MappedOperator, - ValidationSource, - ensure_xcomarg_return_value, - get_mappable_types, -) +from airflow.models.mappedoperator import MappedOperator, ValidationSource, ensure_xcomarg_return_value from airflow.models.pool import Pool from airflow.models.xcom_arg import XComArg from airflow.typing_compat import ParamSpec, Protocol @@ -100,7 +96,7 @@ def _validate_arg_names(self, func: ValidationSource, kwargs: dict[str, Any]) -> kwargs_left = kwargs.copy() for arg_name in self._mappable_function_argument_names: value = kwargs_left.pop(arg_name, NOTSET) - if func != "expand" or value is NOTSET or isinstance(value, get_mappable_types()): + if func != "expand" or value is NOTSET or is_mappable(value): continue tname = type(value).__name__ raise ValueError(f"expand() got an unexpected type {tname!r} for keyword argument {arg_name!r}") diff --git a/airflow/decorators/task_group.py b/airflow/decorators/task_group.py index 6d19ced3d731a..2aa714be3b6c5 100644 --- a/airflow/decorators/task_group.py +++ b/airflow/decorators/task_group.py @@ -30,23 +30,19 @@ from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, Mapping, Sequence, TypeVar, overload import attr -from sqlalchemy.orm import Session from airflow.decorators.base import ExpandableFactory from airflow.models.expandinput import ( DictOfListsExpandInput, - ExpandInput, ListOfDictsExpandInput, + MappedArgument, OperatorExpandArgument, OperatorExpandKwargsArgument, ) from airflow.models.taskmixin import DAGNode from airflow.models.xcom_arg import XComArg from airflow.typing_compat import ParamSpec -from airflow.utils.context import Context from airflow.utils.helpers import prevent_duplicates -from airflow.utils.mixins import ResolveMixin -from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.task_group import MappedTaskGroup, TaskGroup if TYPE_CHECKING: @@ -58,17 +54,6 @@ task_group_sig = inspect.signature(TaskGroup.__init__) -@attr.define(kw_only=True) -class _MappedArgument(ResolveMixin): - _input: ExpandInput - _key: str - - @provide_session - def resolve(self, context: Context, *, session: Session = NEW_SESSION) -> Any: - data, _ = self._input.resolve(context, session=session) - return data[self._key] - - @attr.define() class _TaskGroupFactory(ExpandableFactory, Generic[FParams, FReturn]): function: Callable[FParams, FReturn] = attr.ib(validator=attr.validators.is_callable()) @@ -146,7 +131,7 @@ def expand(self, **kwargs: OperatorExpandArgument) -> DAGNode: return self._create_task_group( functools.partial(MappedTaskGroup, expand_input=expand_input), **self.partial_kwargs, - **{k: _MappedArgument(input=expand_input, key=k) for k in kwargs}, + **{k: MappedArgument(input=expand_input, key=k) for k in kwargs}, ) def expand_kwargs(self, kwargs: OperatorExpandKwargsArgument) -> DAGNode: @@ -175,7 +160,7 @@ def expand_kwargs(self, kwargs: OperatorExpandKwargsArgument) -> DAGNode: return self._create_task_group( functools.partial(MappedTaskGroup, expand_input=expand_input), **self.partial_kwargs, - **{k: _MappedArgument(input=expand_input, key=k) for k in map_kwargs}, + **{k: MappedArgument(input=expand_input, key=k) for k in map_kwargs}, ) diff --git a/airflow/models/expandinput.py b/airflow/models/expandinput.py index 50094933bc597..b817da38dfcde 100644 --- a/airflow/models/expandinput.py +++ b/airflow/models/expandinput.py @@ -22,8 +22,12 @@ import operator from typing import TYPE_CHECKING, Any, Dict, Iterable, Mapping, NamedTuple, Sequence, Sized, Union -from airflow.compat.functools import cache +import attr + +from airflow.typing_compat import TypeGuard from airflow.utils.context import Context +from airflow.utils.mixins import ResolveMixin +from airflow.utils.session import NEW_SESSION, provide_session if TYPE_CHECKING: from sqlalchemy.orm import Session @@ -34,19 +38,54 @@ # Each keyword argument to expand() can be an XComArg, sequence, or dict (not # any mapping since we need the value to be ordered). -OperatorExpandArgument = Union["XComArg", Sequence, Dict[str, Any]] +OperatorExpandArgument = Union["MappedArgument", "XComArg", Sequence, Dict[str, Any]] # The single argument of expand_kwargs() can be an XComArg, or a list with each # element being either an XComArg or a dict. OperatorExpandKwargsArgument = Union["XComArg", Sequence[Union["XComArg", Mapping[str, Any]]]] -# For isinstance() check. -@cache -def get_mappable_types() -> tuple[type, ...]: +@attr.define(kw_only=True) +class MappedArgument(ResolveMixin): + """Stand-in stub for task-group-mapping arguments. + + This is very similar to an XComArg, but resolved differently. Declared here + (instead of in the task group module) to avoid import cycles. + """ + + _input: ExpandInput + _key: str + + def get_task_map_length(self, run_id: str, *, session: Session) -> int | None: + # TODO (AIP-42): Implement run-time task map length inspection. + # This simply marks the value as un-expandable at parse-time. + return None + + @provide_session + def resolve(self, context: Context, *, session: Session = NEW_SESSION) -> Any: + data, _ = self._input.resolve(context, session=session) + return data[self._key] + + +# To replace tedious isinstance() checks. +def is_mappable(v: Any) -> TypeGuard[OperatorExpandArgument]: + from airflow.models.xcom_arg import XComArg + + return isinstance(v, (MappedArgument, XComArg, Mapping, Sequence)) and not isinstance(v, str) + + +# To replace tedious isinstance() checks. +def _is_parse_time_mappable(v: OperatorExpandArgument) -> TypeGuard[Mapping | Sequence]: + from airflow.models.xcom_arg import XComArg + + return not isinstance(v, (MappedArgument, XComArg)) + + +# To replace tedious isinstance() checks. +def _needs_run_time_resolution(v: OperatorExpandArgument) -> TypeGuard[MappedArgument | XComArg]: from airflow.models.xcom_arg import XComArg - return (XComArg, list, tuple, dict) + return isinstance(v, (MappedArgument, XComArg)) class NotFullyPopulated(RuntimeError): @@ -74,9 +113,7 @@ class DictOfListsExpandInput(NamedTuple): def _iter_parse_time_resolved_kwargs(self) -> Iterable[tuple[str, Sized]]: """Generate kwargs with values available on parse-time.""" - from airflow.models.xcom_arg import XComArg - - return ((k, v) for k, v in self.value.items() if not isinstance(v, XComArg)) + return ((k, v) for k, v in self.value.items() if _is_parse_time_mappable(v)) def get_parse_time_mapped_ti_count(self) -> int: if not self.value: @@ -93,14 +130,18 @@ def _get_map_lengths(self, run_id: str, *, session: Session) -> dict[str, int]: If any arguments are not known right now (upstream task not finished), they will not be present in the dict. """ - from airflow.models.xcom_arg import XComArg - # TODO: This initiates one database call for each XComArg. Would it be # more efficient to do one single db call and unpack the value here? - map_lengths_iterator = ( - (k, (v.get_task_map_length(run_id, session=session) if isinstance(v, XComArg) else len(v))) - for k, v in self.value.items() - ) + def _get_length(v: OperatorExpandArgument) -> int | None: + if _needs_run_time_resolution(v): + return v.get_task_map_length(run_id, session=session) + # Unfortunately a user-defined TypeGuard cannot apply negative type + # narrowing. https://github.com/python/typing/discussions/1013 + if TYPE_CHECKING: + assert isinstance(v, Sized) + return len(v) + + map_lengths_iterator = ((k, _get_length(v)) for k, v in self.value.items()) map_lengths = {k: v for k, v in map_lengths_iterator if v is not None} if len(map_lengths) < len(self.value): @@ -114,9 +155,7 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int: return functools.reduce(operator.mul, (lengths[name] for name in self.value), 1) def _expand_mapped_field(self, key: str, value: Any, context: Context, *, session: Session) -> Any: - from airflow.models.xcom_arg import XComArg - - if isinstance(value, XComArg): + if _needs_run_time_resolution(value): value = value.resolve(context, session=session) map_index = context["ti"].map_index if map_index < 0: diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py index 35f56279e76af..ca6a3340948be 100644 --- a/airflow/models/mappedoperator.py +++ b/airflow/models/mappedoperator.py @@ -54,7 +54,7 @@ NotFullyPopulated, OperatorExpandArgument, OperatorExpandKwargsArgument, - get_mappable_types, + is_mappable, ) from airflow.models.param import ParamsDict from airflow.models.pool import Pool @@ -97,7 +97,7 @@ def validate_mapping_kwargs(op: type[BaseOperator], func: ValidationSource, valu continue if value is NOTSET: continue - if isinstance(value, get_mappable_types()): + if is_mappable(value): continue type_name = type(value).__name__ error = f"{op.__name__}.expand() got an unexpected type {type_name!r} for keyword argument {name}" diff --git a/airflow/typing_compat.py b/airflow/typing_compat.py index ee9cdab7a52f4..5f2321962a9f8 100644 --- a/airflow/typing_compat.py +++ b/airflow/typing_compat.py @@ -26,6 +26,7 @@ "ParamSpec", "Protocol", "TypedDict", + "TypeGuard", "runtime_checkable", ] @@ -43,6 +44,6 @@ from typing_extensions import Literal if sys.version_info >= (3, 10): - from typing import ParamSpec + from typing import ParamSpec, TypeGuard else: - from typing_extensions import ParamSpec + from typing_extensions import ParamSpec, TypeGuard diff --git a/tests/decorators/test_task_group.py b/tests/decorators/test_task_group.py index 3e9d143ffb717..38b54cf1c52f0 100644 --- a/tests/decorators/test_task_group.py +++ b/tests/decorators/test_task_group.py @@ -21,8 +21,7 @@ import pytest from airflow.decorators import dag, task_group -from airflow.decorators.task_group import _MappedArgument -from airflow.models.expandinput import DictOfListsExpandInput, ListOfDictsExpandInput +from airflow.models.expandinput import DictOfListsExpandInput, ListOfDictsExpandInput, MappedArgument from airflow.utils.task_group import MappedTaskGroup @@ -150,7 +149,7 @@ def tg(a, b): assert isinstance(tg, MappedTaskGroup) assert tg._expand_input == DictOfListsExpandInput({"b": ["x", "y"]}) - assert saved == {"a": 1, "b": _MappedArgument(input=tg._expand_input, key="b")} + assert saved == {"a": 1, "b": MappedArgument(input=tg._expand_input, key="b")} def test_expand_kwargs_no_wildcard(): @@ -186,4 +185,4 @@ def tg(a, b): assert isinstance(tg, MappedTaskGroup) assert tg._expand_input == ListOfDictsExpandInput([{"b": "x"}, {"b": None}]) - assert saved == {"a": 1, "b": _MappedArgument(input=tg._expand_input, key="b")} + assert saved == {"a": 1, "b": MappedArgument(input=tg._expand_input, key="b")} diff --git a/tests/models/test_dagrun.py b/tests/models/test_dagrun.py index 5818df321bf52..6295dbcf90c16 100644 --- a/tests/models/test_dagrun.py +++ b/tests/models/test_dagrun.py @@ -28,7 +28,7 @@ from airflow import settings from airflow.callbacks.callback_requests import DagCallbackRequest -from airflow.decorators import task +from airflow.decorators import task, task_group from airflow.models import DAG, DagBag, DagModel, DagRun, TaskInstance as TI, clear_task_instances from airflow.models.baseoperator import BaseOperator from airflow.models.taskmap import TaskMap @@ -1670,7 +1670,7 @@ def task_2(arg2): @pytest.mark.need_serialized_dag -def test_mapped_mixed__literal_not_expanded_at_create(dag_maker, session): +def test_mapped_mixed_literal_not_expanded_at_create(dag_maker, session): literal = [1, 2, 3, 4] with dag_maker(session=session): task = BaseOperator(task_id="task_1") @@ -1690,6 +1690,43 @@ def test_mapped_mixed__literal_not_expanded_at_create(dag_maker, session): assert query.all() == [(-1, None)] +def test_mapped_task_group_expands_at_create(dag_maker, session): + literal = [[1, 2], [3, 4]] + + with dag_maker(session=session): + + @task_group + def tg(x): + # Normal operator in mapped task group, expands to 2 tis. + MockOperator(task_id="t1") + # Mapped operator expands *again* against mapped task group arguments to 4 tis. + MockOperator.partial(task_id="t2").expand(arg1=literal) + # Normal operator referencing mapped task group arguments does not further expand, only 2 tis. + MockOperator(task_id="t3", arg1=x) + # It can expand *again* (since each item in x is a list) but this is not done at parse time. + MockOperator.partial(task_id="t4").expand(arg1=x) + + tg.expand(x=literal) + + dr = dag_maker.create_dagrun() + query = ( + session.query(TI.task_id, TI.map_index, TI.state) + .filter_by(dag_id=dr.dag_id, run_id=dr.run_id) + .order_by(TI.task_id, TI.map_index) + ) + assert query.all() == [ + ("tg.t1", 0, None), + ("tg.t1", 1, None), + ("tg.t2", 0, None), + ("tg.t2", 1, None), + ("tg.t2", 2, None), + ("tg.t2", 3, None), + ("tg.t3", 0, None), + ("tg.t3", 1, None), + ("tg.t4", -1, None), + ] + + def test_ti_scheduling_mapped_zero_length(dag_maker, session): with dag_maker(session=session): task = BaseOperator(task_id="task_1")