fixup! Offer task loops as a public authoring API
diff --git a/airflow-core/docs/authoring-and-scheduling/examples/example_task_loops.py b/airflow-core/docs/authoring-and-scheduling/examples/example_task_loops.py deleted file mode 100644 index 89ec482..0000000 --- a/airflow-core/docs/authoring-and-scheduling/examples/example_task_loops.py +++ /dev/null
@@ -1,110 +0,0 @@ -# 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. -"""Repeat task groups with fixed counts, runtime conditions, and mapped tasks.""" - -from __future__ import annotations - -# [START refine_estimate] -from airflow.sdk import dag, task, task_group - - -@dag(schedule=None, catchup=False, tags=["example"]) -def refine_estimate(): - @task_group - def refine(): - @task - def improve(*, loop): - estimate = 1.0 if loop.index == 0 else loop.previous["estimate"] - return (estimate + 2.0 / estimate) / 2.0 - - @task - def evaluate(estimate): - return {"estimate": estimate, "error": abs(estimate * estimate - 2.0)} - - evaluate(improve()) - - def accurate_enough(*, loop): - return loop.result["error"] < 0.000001 - - @task - def finished(): - print("Refinement finished.") - - refinement = refine.loop(max_iterations=10, until=accurate_enough) - refinement >> finished() - - -refine_estimate() -# [END refine_estimate] - - -# [START fixed_loop] -@dag(schedule=None, catchup=False, tags=["example"]) -def fixed_task_loop(): - @task_group - def accumulate(): - @task - def increment(*, loop): - return (0 if loop.index == 0 else loop.previous) + 1 - - increment() - - accumulate.loop(max_iterations=3) - - -fixed_task_loop() -# [END fixed_loop] - - -# [START mapped_loop] -@dag(schedule=None, catchup=False, tags=["example"]) -def mapped_task_loop(): - @task_group - def process_batch(): - @task - def process(value, *, loop, ti): - print(f"Iteration {loop.index}, mapped position {ti.map_index}") - return value + loop.index - - process.expand(value=[1, 2]) - - def batch_ready(*, loop): - return min(loop.result) >= 2 - - process_batch.loop(max_iterations=3, until=batch_ready) - - -mapped_task_loop() -# [END mapped_loop] - - -# [START partial_override_loop] -@dag(schedule=None, catchup=False, tags=["example"]) -def partial_override_loop(): - @task_group - def accumulate(increment_by): - @task - def increment(increment_by, *, loop): - return (0 if loop.index == 0 else loop.previous) + increment_by - - increment(increment_by) - - accumulate.override(group_id="accumulate_more").partial(increment_by=5).loop(max_iterations=3) - - -partial_override_loop() -# [END partial_override_loop]
diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst b/airflow-core/docs/authoring-and-scheduling/loops.rst index 8a411f9..1262883 100644 --- a/airflow-core/docs/authoring-and-scheduling/loops.rst +++ b/airflow-core/docs/authoring-and-scheduling/loops.rst
@@ -58,7 +58,7 @@ This example improves an estimate of the square root of two until the error is small enough: -.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py +.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py :start-after: [START refine_estimate] :end-before: [END refine_estimate] @@ -81,7 +81,7 @@ three iterations, carrying results between them. Reaching the cap completes a fixed-count loop successfully. Its gate is named ``__loop_gate`` within the group. -.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py +.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py :start-after: [START fixed_loop] :end-before: [END fixed_loop] @@ -89,7 +89,7 @@ calling ``.loop()``. Use ``.override()`` to configure the group, for example to give another loop a different ``group_id``: -.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py +.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py :start-after: [START partial_override_loop] :end-before: [END partial_override_loop] @@ -187,7 +187,7 @@ under the gate's ``all_success`` rule, skips the gate; no next iteration is created. -.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py +.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py :start-after: [START mapped_loop] :end-before: [END mapped_loop]
diff --git a/airflow-core/src/airflow/example_dags/example_task_loops.py b/airflow-core/src/airflow/example_dags/example_task_loops.py index 9925a63..89ec482 100644 --- a/airflow-core/src/airflow/example_dags/example_task_loops.py +++ b/airflow-core/src/airflow/example_dags/example_task_loops.py
@@ -28,8 +28,7 @@ def refine(): @task def improve(*, loop): - previous = loop.previous - estimate = 1.0 if previous is None else previous["estimate"] + estimate = 1.0 if loop.index == 0 else loop.previous["estimate"] return (estimate + 2.0 / estimate) / 2.0 @task @@ -60,8 +59,7 @@ def accumulate(): @task def increment(*, loop): - previous = loop.previous - return (0 if previous is None else previous) + 1 + return (0 if loop.index == 0 else loop.previous) + 1 increment() @@ -92,3 +90,21 @@ mapped_task_loop() # [END mapped_loop] + + +# [START partial_override_loop] +@dag(schedule=None, catchup=False, tags=["example"]) +def partial_override_loop(): + @task_group + def accumulate(increment_by): + @task + def increment(increment_by, *, loop): + return (0 if loop.index == 0 else loop.previous) + increment_by + + increment(increment_by) + + accumulate.override(group_id="accumulate_more").partial(increment_by=5).loop(max_iterations=3) + + +partial_override_loop() +# [END partial_override_loop]
diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index fccd20b..db0069f 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -22,7 +22,6 @@ from collections.abc import Mapping from contextlib import contextmanager from functools import partial, reduce -from pathlib import Path from typing import TYPE_CHECKING from unittest import mock from unittest.mock import ANY, call @@ -99,9 +98,10 @@ from tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.config import conf_vars from tests_common.test_utils.dag import sync_dag_to_db -from tests_common.test_utils.db import clear_db_dags, clear_db_runs, clear_db_xcom +from tests_common.test_utils.db import clear_db_dags, clear_db_runs, clear_db_serialized_dags, clear_db_xcom from tests_common.test_utils.mapping import expand_mapped_task, push_mapped_length from tests_common.test_utils.mock_operators import MockOperator +from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH from tests_common.test_utils.taskinstance import create_task_instance, run_task_instance from unit.models import DEFAULT_DATE as _DEFAULT_DATE from unit.plugins.priority_weight_strategy import DecreasingPriorityStrategy, TestPriorityWeightStrategyPlugin @@ -206,7 +206,7 @@ session.flush() with dag_maker(serialized=True, session=session): - create_loop(body, max_iterations=3, add_member=True) + body.partial(add_member=True).loop(max_iterations=3) dr.dag = dag_maker.serialized_dag version_id = DagVersion.get_latest_version(dr.dag_id, session=session).id @@ -303,16 +303,23 @@ yield clear_db_runs() clear_db_xcom() + clear_db_dags() + clear_db_serialized_dags() @pytest.mark.usefixtures("clear_loop_example_runs") @pytest.mark.parametrize( ("dag_id", "expected_iterations"), - [("refine_estimate", 4), ("fixed_task_loop", 3), ("mapped_task_loop", 2)], + [ + ("refine_estimate", 4), + ("fixed_task_loop", 3), + ("mapped_task_loop", 2), + ("partial_override_loop", 3), + ], ) def test_documented_task_loop_examples_complete(session, dag_id, expected_iterations): bag = DagBag( - dag_folder=str(Path(__file__).parents[3] / "src/airflow/example_dags/example_task_loops.py"), + dag_folder=str(AIRFLOW_CORE_SOURCES_PATH / "airflow" / "example_dags" / "example_task_loops.py"), ) assert bag.import_errors == {} dag = bag.get_dag(dag_id) @@ -320,11 +327,7 @@ dr = dag.test() assert dr.state == DagRunState.SUCCESS - gates = [ - ti - for ti in dr.get_task_instances(session=session) - if ti.operator == "LoopGateOperator" and ti.working_set - ] + gates = [ti for ti in dr.get_task_instances(session=session) if ti.operator == "LoopGateOperator"] assert sorted((ti.region_index, ti.state) for ti in gates) == [ (index, State.SUCCESS) for index in range(expected_iterations) ]