Clean up SDK references in airflow.models.expandinput (#59815)

Author: uranusjr

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -64,11 +64,11 @@
 from airflow.models.asset import AssetActive
 from airflow.models.dag import DagModel
 from airflow.models.dagrun import DagRun as DR
+from airflow.models.expandinput import NotFullyPopulated
 from airflow.models.taskinstance import TaskInstance as TI, _stop_remaining_tasks
 from airflow.models.taskreschedule import TaskReschedule
 from airflow.models.trigger import Trigger
 from airflow.models.xcom import XComModel
-from airflow.sdk.definitions._internal.expandinput import NotFullyPopulated
 from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey
 from airflow.serialization.definitions.dag import SerializedDAG
 from airflow.task.trigger_rule import TriggerRule

airflow-core/src/airflow/models/expandinput.py

@@ -24,45 +24,83 @@
 
 import attrs
 
-from airflow.sdk.definitions._internal.expandinput import (
-    DictOfListsExpandInput,
-    ListOfDictsExpandInput,
-    MappedArgument,
-    NotFullyPopulated,
-    OperatorExpandArgument,
-    OperatorExpandKwargsArgument,
-    is_mappable,
-)
-
 if TYPE_CHECKING:
-    from typing import TypeGuard
+    from collections.abc import Mapping, Sequence
+    from typing import TypeAlias
 
     from sqlalchemy.orm import Session
+    from typing_extensions import TypeIs
 
     from airflow.serialization.definitions.mappedoperator import Operator
     from airflow.serialization.definitions.xcom_arg import SchedulerXComArg
 
+    ExpandArgument: TypeAlias = "SchedulerMappedArgument" | SchedulerXComArg | Sequence | Mapping[str, Any]
+    ExpandKwargsArgument: TypeAlias = SchedulerXComArg | Sequence[SchedulerXComArg | Mapping[str, Any]]
+
 
 __all__ = [
-    "DictOfListsExpandInput",
-    "ListOfDictsExpandInput",
-    "MappedArgument",
     "NotFullyPopulated",
-    "OperatorExpandArgument",
-    "OperatorExpandKwargsArgument",
-    "is_mappable",
+    "SchedulerMappedArgument",
+    "SchedulerDictOfListsExpandInput",
+    "SchedulerListOfDictsExpandInput",
 ]
 
 
-def _needs_run_time_resolution(v: OperatorExpandArgument) -> TypeGuard[MappedArgument | SchedulerXComArg]:
+class NotFullyPopulated(RuntimeError):
+    """
+    Raise when mapped length cannot be calculated due to incomplete metadata.
+
+    This is generally due to not all upstream tasks have been completed (or in
+    parse-time length calculations, when any upstream has runtime dependencies
+    on mapped length) when the function is called.
+    """
+
+    def __init__(self, missing: set[str]) -> None:
+        self.missing = missing
+
+    def __str__(self) -> str:
+        keys = ", ".join(repr(k) for k in sorted(self.missing))
+        return f"Failed to populate all mapping metadata; missing: {keys}"
+
+
+def _needs_run_time_resolution(v: ExpandArgument) -> TypeIs[SchedulerMappedArgument | SchedulerXComArg]:
     from airflow.serialization.definitions.xcom_arg import SchedulerXComArg
 
-    return isinstance(v, (MappedArgument, SchedulerXComArg))
+    return isinstance(v, (SchedulerMappedArgument, SchedulerXComArg))
+
+
+@attrs.define(kw_only=True)
+class SchedulerMappedArgument:
+    """
+    Stand-in stub for task-group-mapping arguments.
+
+    This corresponds on SDK's ``MappedArgument``, which is created when
+    dynamically mapping a task group, and an argument used to dynamic-map is
+    passed into a task inside the group.
+
+    This value is not currently used anywhere in the scheduler since nested
+    dynamic mapping is not supported (i.e. using this value to further expand
+    an operator inside a mapped task group), but this is implemented so the
+    value is displayed better in the UI.
+    """
+
+    _input: SchedulerExpandInput = attrs.field()
+    _key: str
+
+    def iter_references(self) -> Iterable[tuple[Operator, str]]:
+        yield from self._input.iter_references()
 
 
 @attrs.define
 class SchedulerDictOfListsExpandInput:
-    value: dict
+    """
+    Serialized storage of a mapped operator's mapped kwargs.
+
+    This corresponds to SDK's ``DictOfListsExpandInput``, which was created by
+    calling ``expand(**kwargs)`` on an operator type.
+    """
+
+    value: dict[str, ExpandArgument]
 
     EXPAND_INPUT_TYPE: ClassVar[str] = "dict-of-lists"
 
@@ -90,7 +128,7 @@ def _get_map_lengths(self, run_id: str, *, session: Session) -> dict[str, int]:
 
         # 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?
-        def _get_length(v: OperatorExpandArgument) -> int | None:
+        def _get_length(v: ExpandArgument) -> int | None:
             if isinstance(v, SchedulerXComArg):
                 return get_task_map_length(v, run_id, session=session)
 
@@ -123,7 +161,14 @@ def iter_references(self) -> Iterable[tuple[Operator, str]]:
 
 @attrs.define
 class SchedulerListOfDictsExpandInput:
-    value: list
+    """
+    Serialized storage of a mapped operator's mapped kwargs.
+
+    This corresponds to SDK's ``ListOfDictsExpandInput``, which was created by
+    calling ``expand_kwargs(xcom_arg)`` on an operator type.
+    """
+
+    value: ExpandKwargsArgument
 
     EXPAND_INPUT_TYPE: ClassVar[str] = "list-of-dicts"
 

airflow-core/src/airflow/serialization/encoders.py

@@ -59,12 +59,19 @@
 if TYPE_CHECKING:
     from dateutil.relativedelta import relativedelta
 
+    from airflow.sdk.definitions._internal.expandinput import ExpandInput
     from airflow.sdk.definitions.asset import BaseAsset
     from airflow.triggers.base import BaseEventTrigger
 
     T = TypeVar("T")
 
 
+def encode_expand_input(var: ExpandInput) -> dict[str, Any]:
+    from airflow.serialization.serialized_objects import BaseSerialization
+
+    return {"type": var.EXPAND_INPUT_TYPE, "value": BaseSerialization.serialize(var.value)}
+
+
 def encode_relativedelta(var: relativedelta) -> dict[str, Any]:
     """Encode a relativedelta object."""
     encoded = {k: v for k, v in var.__dict__.items() if not k.startswith("_") and v}

airflow-core/src/airflow/serialization/enums.py

@@ -69,3 +69,4 @@ class DagAttributeTypes(str, Enum):
     DAG_CALLBACK_REQUEST = "dag_callback_request"
     TASK_INSTANCE_KEY = "task_instance_key"
     DEADLINE_ALERT = "deadline_alert"
+    MAPPED_ARGUMENT = "mapped_argument"

airflow-core/src/airflow/serialization/serialized_objects.py

@@ -46,10 +46,11 @@
 from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest
 from airflow.exceptions import AirflowException, DeserializationError, SerializationError
 from airflow.models.connection import Connection
-from airflow.models.expandinput import create_expand_input
+from airflow.models.expandinput import SchedulerMappedArgument, create_expand_input
 from airflow.models.taskinstancekey import TaskInstanceKey
 from airflow.sdk import DAG, Asset, AssetAlias, BaseOperator, XComArg
 from airflow.sdk.bases.operator import OPERATOR_DEFAULTS  # TODO: Copy this into the scheduler?
+from airflow.sdk.definitions._internal.expandinput import MappedArgument
 from airflow.sdk.definitions.asset import (
     AssetAliasEvent,
     AssetAliasUniqueKey,
@@ -81,6 +82,7 @@
 from airflow.serialization.encoders import (
     coerce_to_core_timetable,
     encode_asset_like,
+    encode_expand_input,
     encode_relativedelta,
     encode_timetable,
     encode_timezone,
@@ -654,6 +656,9 @@ def serialize(
             return cls._encode(var.to_json(), type_=DAT.TASK_CALLBACK_REQUEST)
         elif isinstance(var, DagCallbackRequest):
             return cls._encode(var.to_json(), type_=DAT.DAG_CALLBACK_REQUEST)
+        elif isinstance(var, MappedArgument):
+            data = {"input": encode_expand_input(var._input), "key": var._key}
+            return cls._encode(data, type_=DAT.MAPPED_ARGUMENT)
         elif var.__class__ == Context:
             d = {}
             for k, v in var.items():
@@ -763,6 +768,9 @@ def deserialize(cls, encoded_var: Any) -> Any:
             return DagCallbackRequest.from_json(var)
         elif type_ == DAT.TASK_INSTANCE_KEY:
             return TaskInstanceKey(**var)
+        elif type_ == DAT.MAPPED_ARGUMENT:
+            expand_input = create_expand_input(var["input"]["type"], var["input"]["value"])
+            return SchedulerMappedArgument(input=expand_input, key=var["key"])
         elif type_ == DAT.ARG_NOT_SET:
             from airflow.serialization.definitions.notset import NOTSET
 
@@ -1033,10 +1041,7 @@ def serialize_mapped_operator(cls, op: MappedOperator) -> dict[str, Any]:
         expansion_kwargs = op._get_specified_expand_input()
         if TYPE_CHECKING:  # Let Mypy check the input type for us!
             _ExpandInputRef.validate_expand_input_value(expansion_kwargs.value)
-        serialized_op[op._expand_input_attr] = {
-            "type": type(expansion_kwargs).EXPAND_INPUT_TYPE,
-            "value": cls.serialize(expansion_kwargs.value),
-        }
+        serialized_op[op._expand_input_attr] = encode_expand_input(expansion_kwargs)
 
         if op.partial_kwargs:
             serialized_op["partial_kwargs"] = {}
@@ -2179,11 +2184,7 @@ def serialize_task_group(cls, task_group: TaskGroup) -> dict[str, Any] | None:
         }
 
         if isinstance(task_group, MappedTaskGroup):
-            expand_input = task_group._expand_input
-            encoded["expand_input"] = {
-                "type": expand_input.EXPAND_INPUT_TYPE,
-                "value": cls.serialize(expand_input.value),
-            }
+            encoded["expand_input"] = encode_expand_input(task_group._expand_input)
             encoded["is_mapped"] = True
 
         return encoded

airflow-core/tests/unit/models/test_dagrun.py

@@ -1456,7 +1456,7 @@ def task_2(arg2): ...
         task_2.expand(arg2=[1, 2])
 
     # Update it to use the new serialized DAG
-    dr.dag = dag_maker.dag
+    dr.dag = dag_maker.serialized_model.dag
     dag_version_id = DagVersion.get_latest_version(dag_id=dr.dag_id, session=session).id
     dr.verify_integrity(dag_version_id=dag_version_id, session=session)
 
@@ -1522,10 +1522,10 @@ def task_2(arg2): ...
     ]
 
     # Now "increase" the length of literal
-    with dag_maker(session=session, serialized=True) as dag:
+    with dag_maker(session=session, serialized=True):
         task_2.expand(arg2=[1, 2, 3, 4, 5])
 
-    dr.dag = dag
+    dr.dag = dag_maker.serialized_model.dag
     # Every mapped task is revised at task_instance_scheduling_decision
     dr.task_instance_scheduling_decisions()
 
@@ -1565,7 +1565,7 @@ def task_2(arg2): ...
     with dag_maker(session=session):
         task_2.expand(arg2=[1, 2])
 
-    dr.dag = dag_maker.dag
+    dr.dag = dag_maker.serialized_model.dag
     # Since we change the literal on the dag file itself, the dag_hash will
     # change which will have the scheduler verify the dr integrity
     dag_version_id = DagVersion.get_latest_version(dag_id=dr.dag_id, session=session).id

providers/standard/tests/unit/standard/decorators/test_python.py

@@ -48,7 +48,7 @@
     from airflow.decorators.base import DecoratedMappedOperator  # type: ignore[no-redef]
     from airflow.models.baseoperator import BaseOperator  # type: ignore[no-redef]
     from airflow.models.dag import DAG  # type: ignore[assignment,no-redef]
-    from airflow.models.expandinput import DictOfListsExpandInput
+    from airflow.models.expandinput import DictOfListsExpandInput  # type: ignore[attr-defined,no-redef]
     from airflow.models.xcom_arg import XComArg  # type: ignore[no-redef]
     from airflow.utils.task_group import TaskGroup  # type: ignore[no-redef]
 

task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py

@@ -28,8 +28,6 @@
 )
 from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias
 
-import methodtools
-
 from airflow.sdk import TriggerRule, WeightRule
 from airflow.sdk.configuration import conf
 from airflow.sdk.definitions._internal.mixins import DependencyMixin
@@ -83,10 +81,6 @@
 log = logging.getLogger(__name__)
 
 
-class NotMapped(Exception):
-    """Raise if a task is neither mapped nor has any parent mapped groups."""
-
-
 class AbstractOperator(Templater, DAGNode):
     """
     Common implementation for operators, including unmapped and mapped.
@@ -408,21 +402,3 @@ def get_needs_expansion(self) -> bool:
             else:
                 self._needs_expansion = False
         return self._needs_expansion
-
-    @methodtools.lru_cache(maxsize=None)
-    def get_parse_time_mapped_ti_count(self) -> int:
-        """
-        Return the 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.
-        """
-        group = self.get_closest_mapped_task_group()
-        if group is None:
-            raise NotMapped()
-        return group.get_parse_time_mapped_ti_count()

task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py

@@ -17,8 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import functools
-import operator
 from collections.abc import Iterable, Mapping, Sequence, Sized
 from typing import TYPE_CHECKING, Any, ClassVar, Union
 
@@ -43,12 +41,13 @@
 OperatorExpandKwargsArgument = Union["XComArg", Sequence[Union["XComArg", Mapping[str, Any]]]]
 
 
-class NotFullyPopulated(RuntimeError):
+class _NotFullyPopulated(RuntimeError):
     """
-    Raise when ``get_map_lengths`` cannot populate all mapping metadata.
+    Raise when an expand input cannot be resolved due to incomplete metadata.
 
-    This is generally due to not all upstream tasks have finished when the
-    function is called.
+    This generally should not happen. The scheduler should have made sure that
+    a not-yet-ready-to-expand task should not be executed. In the off chance
+    this gets raised, it will fail the task instance.
     """
 
     def __init__(self, missing: set[str]) -> None:
@@ -123,15 +122,6 @@ def _iter_parse_time_resolved_kwargs(self) -> Iterable[tuple[str, Sized]]:
         """Generate kwargs with values available on parse-time."""
         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:
-            return 0
-        literal_values = [len(v) for _, v in self._iter_parse_time_resolved_kwargs()]
-        if len(literal_values) != len(self.value):
-            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, resolved_vals: dict[str, Sized], upstream_map_indexes: dict[str, int]
     ) -> dict[str, int]:
@@ -160,7 +150,7 @@ def _get_length(k: str, v: OperatorExpandArgument) -> int | None:
             k: res for k, v in self.value.items() if v is not None if (res := _get_length(k, v)) is not None
         }
         if len(map_lengths) < len(self.value):
-            raise NotFullyPopulated(set(self.value).difference(map_lengths))
+            raise _NotFullyPopulated(set(self.value).difference(map_lengths))
         return map_lengths
 
     def _expand_mapped_field(self, key: str, value: Any, map_index: int, all_lengths: dict[str, int]) -> Any:
@@ -236,11 +226,6 @@ class ListOfDictsExpandInput(ResolveMixin):
 
     EXPAND_INPUT_TYPE: ClassVar[str] = "list-of-dicts"
 
-    def get_parse_time_mapped_ti_count(self) -> int:
-        if isinstance(self.value, Sized):
-            return len(self.value)
-        raise NotFullyPopulated({"expand_kwargs() argument"})
-
     def iter_references(self) -> Iterable[tuple[Operator, str]]:
         from airflow.sdk.definitions.xcom_arg import XComArg