AIP-104: Task Iteration and Dynamic Task Batching

.github/CODEOWNERS

@@ -162,6 +162,9 @@ Dockerfile.ci @potiuk @ashb @gopidesupavan @amoghrajesh @jscheffl @bugraoz93 @ja
 # AIP-72 - Task SDK
 # Python SDK
 /task-sdk/ @ashb @kaxil @amoghrajesh
+/task-sdk/**/batchedoperator.py @dabla
+/task-sdk/**/executor.py @dabla
+/task-sdk/**/iterableoperator.py @dabla
 
 # AIP-108 - Coordinators
 /task-sdk/src/airflow/sdk/coordinators/ @jason810496 @uranusjr

airflow-core/src/airflow/serialization/definitions/mappedoperator.py

@@ -338,6 +338,10 @@ def on_failure_fail_dagrun(self) -> bool:
     def on_failure_fail_dagrun(self, v) -> None:
         self.partial_kwargs["on_failure_fail_dagrun"] = bool(v)
 
+    @property
+    def batch_size(self) -> int:
+        return self.partial_kwargs.get("batch_size", 0)
+
     @classmethod
     def get_serialized_fields(cls):
         """Fields to extract from JSON-Serialized DAG."""
@@ -463,15 +467,20 @@ def get_parse_time_mapped_ti_count(self) -> int:
         current_count = self._get_specified_expand_input().get_parse_time_mapped_ti_count()
 
         def _get_parent_count() -> int:
-            if (group := self.get_closest_mapped_task_group()) is None:
-                raise NotMapped()
-            return group.get_parse_time_mapped_ti_count()
+            try:
+                if (group := self.get_closest_mapped_task_group()) is None:
+                    raise NotMapped()
+                return group.get_parse_time_mapped_ti_count()
+            except NotMapped:
+                return 1
 
-        try:
-            parent_count = _get_parent_count()
-        except NotMapped:
-            return current_count
-        return parent_count * current_count
+        parent_count = _get_parent_count()
+        mapped_ti_count = parent_count * current_count
+
+        if self.batch_size > 0:
+            if self.batch_size < mapped_ti_count:
+                return self.batch_size
+        return mapped_ti_count
 
     def iter_mapped_dependencies(self) -> Iterator[Operator]:
         """Upstream dependencies that provide XComs used by this task for task mapping."""
@@ -519,6 +528,9 @@ def _(task: SerializedBaseOperator | TaskSDKBaseOperator, run_id: str, *, sessio
 def _(task: SerializedMappedOperator | TaskSDKMappedOperator, run_id: str, *, session: Session) -> int:
     from airflow.serialization.serialized_objects import BaseSerialization, _ExpandInputRef
 
+    if task.batch_size > 0:
+        return task.batch_size
+
     exp_input = task._get_specified_expand_input()
     # TODO (GH-52141): 'task' here should be scheduler-bound and returns scheduler expand input.
     if not hasattr(exp_input, "get_total_map_length"):

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

@@ -46,6 +46,7 @@
 from airflow.exceptions import (
     AirflowException,
     AirflowSkipException,
+    NotMapped,
 )
 from airflow.models.asset import (
     AssetActive,
@@ -3487,6 +3488,74 @@ def show(a, b):
             dag_maker.run_ti(ti.task_id, map_index=ti.map_index, dag_run=dag_run, session=session)
         assert outputs == [(2, 5), (2, 10), (4, 5), (4, 10), (8, 5), (8, 10)]
 
+    def test_iterate_literal_cross_product(self, dag_maker, session):
+        """Test an iterated task with literal cross product args properly."""
+        outputs = []
+
+        with dag_maker(dag_id="product_same_types", session=session, serialized=True) as dag:
+
+            @dag.task
+            def show(a, b):
+                outputs.append((a, b))
+
+            show.iterate(a=[2, 4, 8], b=[5, 10])
+
+        dag_run = dag_maker.create_dagrun()
+
+        show_task = dag.get_task("show")
+        with pytest.raises(NotMapped):
+            show_task.get_parse_time_mapped_ti_count()
+        with pytest.raises(NotMapped):
+            TaskMap.expand_mapped_task(show_task, dag_run.run_id, session=session)
+
+        tis = session.scalars(
+            select(TaskInstance)
+            .where(
+                TaskInstance.dag_id == dag.dag_id,
+                TaskInstance.task_id == "show",
+                TaskInstance.run_id == dag_run.run_id,
+            )
+            .order_by(TaskInstance.map_index)
+        ).all()
+        for ti in tis:
+            ti.refresh_from_task(show_task)
+            dag_maker.run_ti(ti.task_id, map_index=ti.map_index, dag_run=dag_run, session=session)
+        assert outputs == [(2, 5), (2, 10), (4, 5), (4, 10), (8, 5), (8, 10)]
+
+    def test_iterate_literal_cross_product_in_batch(self, dag_maker, session):
+        """Test a batched iterated task with literal cross product args properly."""
+        outputs = []
+
+        with dag_maker(dag_id="product_same_types", session=session, serialized=True) as dag:
+
+            @dag.task
+            def show(a, b):
+                outputs.append((a, b))
+
+            show.batch(size=2).iterate(a=[2, 4, 8], b=[5, 10])
+
+        dag_run = dag_maker.create_dagrun()
+
+        show_task = dag.get_task("show")
+        assert show_task.get_parse_time_mapped_ti_count() == 2
+        mapped_tis, max_map_index = TaskMap.expand_mapped_task(show_task, dag_run.run_id, session=session)
+        assert len(mapped_tis) == 0  # Expanded at parse!
+        assert max_map_index == 1
+
+        tis = session.scalars(
+            select(TaskInstance)
+            .where(
+                TaskInstance.dag_id == dag.dag_id,
+                TaskInstance.task_id == "show",
+                TaskInstance.run_id == dag_run.run_id,
+            )
+            .order_by(TaskInstance.map_index)
+        ).all()
+        for ti in tis:
+            ti.refresh_from_task(show_task)
+            dag_maker.run_ti(ti.task_id, map_index=ti.map_index, dag_run=dag_run, session=session)
+        assert outputs == [(2, 5), (4, 5), (8, 5), (2, 10), (4, 10), (8, 10)]
+
     def test_map_in_group(self, tmp_path: pathlib.Path, dag_maker, session):
         out = tmp_path.joinpath("out")
         out.touch()

airflow-core/tests/unit/serialization/test_dag_serialization.py

@@ -2874,6 +2874,7 @@ def test_operator_expand_serde():
         "template_ext": [".sh", ".bash"],
         "template_fields_renderers": {"bash_command": "bash", "env": "json"},
         "ui_color": "#f0ede4",
+        "_register_with_dag": True,
         "_disallow_kwargs_override": False,
         "_expand_input_attr": "expand_input",
     }
@@ -2917,6 +2918,7 @@ def test_operator_expand_xcomarg_serde():
         },
         "task_id": "task_2",
         "template_fields": ["arg1", "arg2"],
+        "_register_with_dag": True,
         "_disallow_kwargs_override": False,
         "_expand_input_attr": "expand_input",
     }
@@ -2972,6 +2974,7 @@ def test_operator_expand_kwargs_literal_serde(strict):
         },
         "task_id": "task_2",
         "template_fields": ["arg1", "arg2"],
+        "_register_with_dag": True,
         "_disallow_kwargs_override": strict,
         "_expand_input_attr": "expand_input",
     }
@@ -3019,6 +3022,7 @@ def test_operator_expand_kwargs_xcomarg_serde(strict):
         },
         "task_id": "task_2",
         "template_fields": ["arg1", "arg2"],
+        "_register_with_dag": True,
         "_disallow_kwargs_override": strict,
         "_expand_input_attr": "expand_input",
     }
@@ -3127,6 +3131,7 @@ def x(arg1, arg2, arg3):
             "op_args": "py",
             "op_kwargs": "py",
         },
+        "_register_with_dag": True,
         "_disallow_kwargs_override": False,
         "_expand_input_attr": "op_kwargs_expand_input",
         "python_callable_name": "test_taskflow_expand_serde.<locals>.x",
@@ -3220,6 +3225,7 @@ def x(arg1, arg2, arg3):
             "op_args": "py",
             "op_kwargs": "py",
         },
+        "_register_with_dag": True,
         "_disallow_kwargs_override": strict,
         "_expand_input_attr": "op_kwargs_expand_input",
     }
@@ -3331,6 +3337,7 @@ def operator_extra_links(self):
         "partial_kwargs": {
             "retry_delay": {"__type": "timedelta", "__var": 300.0},
         },
+        "_register_with_dag": True,
         "_disallow_kwargs_override": False,
         "_expand_input_attr": "expand_input",
         "_operator_extra_links": {"airflow": "_link_AirflowLink2"},

devel-common/src/tests_common/test_utils/mock_context.py

@@ -20,14 +20,29 @@
 from typing import TYPE_CHECKING, Any
 from unittest import mock
 
+from airflow.models import DagRun
+from airflow.utils.types import DagRunType
+
 from tests_common.test_utils.compat import Context
 from tests_common.test_utils.taskinstance import create_task_instance
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS
+
+if AIRFLOW_V_3_1_PLUS:
+    from airflow.sdk import timezone
+else:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
 
-def mock_context(task) -> Context:
+def generate_run_id() -> str:
+    if AIRFLOW_V_3_0_PLUS:
+        return DagRun.generate_run_id(run_type=DagRunType.MANUAL, run_after=timezone.utcnow())
+    return DagRun.generate_run_id(run_type=DagRunType.MANUAL, execution_date=timezone.utcnow())  # type: ignore[call-arg]
+
+
+def mock_context(task, run_id: str | None = None) -> Context:
     from airflow.models import TaskInstance
     from airflow.utils.session import NEW_SESSION
 
@@ -64,6 +79,6 @@ def xcom_push(self, key: str, value: Any, session: Session = NEW_SESSION, **kwar
             values[key] = value
 
     values["ti"] = create_task_instance(task, dag_version_id=mock.MagicMock(), ti_type=MockedTaskInstance)
+    values["run_id"] = generate_run_id() if run_id is None else run_id
 
-    # See https://github.com/python/mypy/issues/8890 - mypy does not support passing typed dict to TypedDict
     return Context(values)  # type: ignore[misc]

docs/spelling_wordlist.txt

@@ -1260,6 +1260,7 @@ PodManager
 PodSpec
 podSpec
 podspec
+Pokémon
 polars
 poller
 polyfill

task-sdk/.pre-commit-config.yaml

@@ -39,7 +39,9 @@ repos:
           ^src/airflow/sdk/definitions/asset/__init__\.py$|
           ^src/airflow/sdk/definitions/asset/decorators\.py$|
           ^src/airflow/sdk/definitions/taskgroup\.py$|
+          ^src/airflow/sdk/definitions/iterableoperator\.py$|
           ^src/airflow/sdk/definitions/mappedoperator\.py$|
+          ^src/airflow/sdk/definitions/batchedoperator\.py$|
           ^src/airflow/sdk/definitions/deadline\.py$|
           ^src/airflow/sdk/definitions/dag\.py$|
           ^src/airflow/sdk/definitions/_internal/types\.py$|

task-sdk/docs/deferred-vs-async-operators.rst

@@ -205,7 +205,7 @@ concurrently using ``asyncio.gather`` while limiting concurrency with a semaphor
 
 .. note::
 
-   The upcoming *Dynamic Task Iteration* feature will simplify patterns like this.
+   The new :ref:`Dynamic Task Iteration <sdk-dynamic-task-mapping-vs-iteration>`. feature will simplify patterns like this.
    Instead of manually managing concurrency with constructs such as
    ``asyncio.gather`` and ``asyncio.Semaphore``, authors will be able to iterate
    over asynchronous results directly in downstream tasks while still benefiting