From da8c56ca1fdb8ce02556759c3e1fb961af5bd123 Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Fri, 2 Oct 2026 20:36:40 +0100 Subject: [PATCH] Run loop gates and advance iterations atomically The scheduler, the API and the UI must never see every task of a loop finished while its next pass does not exist yet. If we recorded a continuing gate as successful and created the next pass afterwards, a crash between the two would leave recovery unable to tell an unfinished advancement from a completed decision. Gate completion and next-pass creation therefore happen in one transaction under the DagRun lock, which also stops a duplicate completion from advancing the loop twice. The continue-or-stop decision is a reserved XCom stored under the gate's attempt UUID. A retried gate starts without a decision, and a retired attempt cannot advance the loop. An invalid decision moves the gate attempt to retry or failed within the same request; rejecting the request instead would leave the attempt running with nobody to finish it. Gates use normal task dependencies and an explicit loop context. A barrier that waited for every body task would override trigger rules that allow early progress. Workers and dag.test supply the same coordinates, so decisions and previous-iteration reads reach the intended gate and data. Skipping downstream of a task from outside a loop skips every live pass. Skipping only one would let later passes run work the author meant to skip. --- .../docs/authoring-and-scheduling/loops.rst | 8 +- .../execution_api/datamodels/taskinstance.py | 12 + .../execution_api/routes/task_instances.py | 122 +++- .../api_fastapi/execution_api/routes/xcoms.py | 179 ++++- .../execution_api/versions/__init__.py | 6 + .../execution_api/versions/v2026_10_30.py | 11 + .../versions/v2026_10_30_xcom_params.py | 35 + .../example_dags}/example_task_loops.py | 0 .../src/airflow/jobs/scheduler_job_runner.py | 7 +- airflow-core/src/airflow/models/dagrun.py | 144 +++- .../src/airflow/models/dynamic_region.py | 349 ++++++++-- .../src/airflow/models/task_coordinates.py | 415 +++++++++-- .../serialization/definitions/xcom_arg.py | 23 +- .../src/airflow/ti_deps/dep_context.py | 7 +- .../versions/head/test_task_instances.py | 642 +++++++++++++++++- .../execution_api/versions/head/test_xcoms.py | 475 ++++++++++++- .../v2026_10_30/test_task_instances.py | 2 + .../tests/unit/jobs/test_scheduler_job.py | 50 ++ airflow-core/tests/unit/models/test_dagrun.py | 259 ++++++- .../unit/models/test_task_coordinates.py | 241 ++++++- .../tests/unit/models/test_xcom_arg.py | 45 ++ .../serialization/test_dag_serialization.py | 24 + .../ti_deps/deps/test_trigger_rule_dep.py | 74 ++ .../src/tests_common/test_utils/asserts.py | 20 + go-sdk/pkg/execution/genmodels/models.gen.go | 59 +- go-sdk/schema/supervisor-schema.json | 67 ++ .../unit/edge3/worker_api/routes/test_logs.py | 14 +- ...check_template_context_variable_in_sync.py | 2 + task-sdk/src/airflow/sdk/api/client.py | 25 +- .../airflow/sdk/api/datamodels/_generated.py | 13 + task-sdk/src/airflow/sdk/bases/decorator.py | 24 +- task-sdk/src/airflow/sdk/bases/xcom.py | 2 + .../airflow/sdk/definitions/_internal/loop.py | 65 ++ .../airflow/sdk/definitions/_internal/node.py | 15 +- .../src/airflow/sdk/definitions/context.py | 4 +- task-sdk/src/airflow/sdk/definitions/dag.py | 9 +- .../sdk/definitions/decorators/task_group.py | 24 +- .../src/airflow/sdk/definitions/xcom_arg.py | 3 +- .../src/airflow/sdk/execution_time/comms.py | 4 + .../sdk/execution_time/lazy_sequence.py | 4 + .../src/airflow/sdk/execution_time/loop.py | 104 +++ .../sdk/execution_time/request_handlers.py | 11 +- .../sdk/execution_time/schema/schema.json | 67 ++ .../schema/versions/__init__.py | 2 + .../schema/versions/v2026_10_30.py | 17 + .../airflow/sdk/execution_time/task_runner.py | 45 +- .../definitions/_internal/test_loop.py | 231 ++++++- .../definitions/decorators/test_task_group.py | 87 ++- .../task_sdk/definitions/test_xcom_arg.py | 51 +- .../execution_time/schema/test_migrator.py | 39 +- .../task_sdk/execution_time/test_loop.py | 245 +++++++ .../execution_time/test_request_handlers.py | 44 ++ .../execution_time/test_supervisor.py | 28 +- .../execution_time/test_task_runner.py | 64 ++ ts-sdk/schema/supervisor-schema.json | 67 ++ ts-sdk/src/generated/supervisor.ts | 27 + 56 files changed, 4370 insertions(+), 243 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30_xcom_params.py rename airflow-core/{docs/authoring-and-scheduling/examples => src/airflow/example_dags}/example_task_loops.py (100%) create mode 100644 task-sdk/src/airflow/sdk/execution_time/loop.py create mode 100644 task-sdk/tests/task_sdk/execution_time/test_loop.py diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst b/airflow-core/docs/authoring-and-scheduling/loops.rst index 8cd0c76f06d01..b588fdbe42105 100644 --- a/airflow-core/docs/authoring-and-scheduling/loops.rst +++ b/airflow-core/docs/authoring-and-scheduling/loops.rst @@ -68,7 +68,7 @@ iterations. 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] @@ -98,7 +98,7 @@ For a fixed-count loop, omit ``until``. The definition ``accumulate.loop(max_ite 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] @@ -106,7 +106,7 @@ For a task-group function with arguments, supply them with ``.partial()`` before 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] @@ -258,7 +258,7 @@ created. In this example, ``choose_items`` returns the values to map over: ``[1, 2]`` in the first iteration, then each previous result plus one. -.. 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/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index 446aea95b0182..2d158476ada6d 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -427,12 +427,24 @@ def safe_extract_from_orm(cls, data: Any) -> Any: return values +class LoopContext(BaseModel): + """Pinned loop definition and enclosing iteration for a task execution.""" + + node_id: str + index: int = Field(ge=0) + max_iterations: int = Field(gt=0) + terminal_task_id: str + terminal_is_mapped: bool + + class TIRunContext(BaseModel): """Response schema for TaskInstance run context.""" dag_run: DagRun """DAG run information for the task instance.""" + loop: LoopContext | None = None + task_reschedule_count: int = 0 """How many times the task has been rescheduled.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 6c67e87ead9de..99cd52a3e97aa 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -35,7 +35,7 @@ from opentelemetry.trace import StatusCode from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator from pydantic import JsonValue, ValidationError -from sqlalchemy import and_, func, or_, tuple_, update +from sqlalchemy import and_, func, or_, tuple_, union, update from sqlalchemy.engine import CursorResult from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError from sqlalchemy.orm import Session, contains_eager, joinedload @@ -57,6 +57,7 @@ from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( DagRunNoteUpdatePayload, InactiveAssetsResponse, + LoopContext, PreviousTIResponse, PrevSuccessfulDagRunResponse, TaskBreadcrumbsResponse, @@ -95,11 +96,12 @@ from airflow.models.base import ID_LEN from airflow.models.dag import DagModel from airflow.models.dagbag import DBDagBag -from airflow.models.dagrun import DagRun as DR +from airflow.models.dagrun import DagRun as DR, InvalidLoopDecision from airflow.models.dynamic_region import SENTINEL_REGION_ID, AmbiguousProducerError from airflow.models.hitl import HITLDetail from airflow.models.log import Log from airflow.models.task_coordinates import ( + LOOP_GATE_OPERATOR, TaskCoordinateResolver, build_coordinate_filters, get_public_region, @@ -113,7 +115,7 @@ from airflow.state import get_state_backend from airflow.triggers.base import TriggerEvent from airflow.utils.sqlalchemy import get_dialect_name -from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState +from airflow.utils.state import DagRunState, IntermediateTIState, TaskInstanceState, TerminalTIState router = VersionedAPIRouter() @@ -177,6 +179,7 @@ def ti_run( TI.dag_id, TI.run_id, TI.task_id, + TI.region_id, TI.region_index, TI.try_number, TI.max_tries, @@ -336,6 +339,19 @@ def ti_run( should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries), multi_team=conf.getboolean("core", "multi_team"), ) + resolver = TaskCoordinateResolver(dag_bag, session) + if loop_context := resolver.loop_context(ti): + group, index = loop_context + terminal = resolver.get_task( + ti.dag_id, ti.run_id, group.terminal_task_id, dag_version_id=ti.dag_version_id + ) + context.loop = LoopContext( + node_id=group.node_id, + index=index, + max_iterations=group.max_iterations, + terminal_task_id=group.terminal_task_id, + terminal_is_mapped=terminal.get_needs_expansion(), + ) # Only set for lang-SDK (foreign-runtime) tasks with a captured TaskFlow arg # spec; the route excludes unset fields, keeping regular responses lean. @@ -435,6 +451,46 @@ def ti_update_state( raise HTTPException(status_code=409, detail={"reason": "invalid_state"}) return Response(status_code=status.HTTP_204_NO_CONTENT) + loop_group = None + loop_gate = None + if isinstance(ti_patch_payload, (TISuccessStatePayload, TITerminalStatePayload)): + gate_run = session.execute( + select(TI.dag_id, TI.run_id).where( + TI.id == task_instance_id, + TI.working_set.is_(True), + TI.operator == LOOP_GATE_OPERATOR, + ) + ).one_or_none() + if gate_run is not None: + session.execute( + select(DR).where(DR.dag_id == gate_run.dag_id, DR.run_id == gate_run.run_id).with_for_update() + ).scalar_one() + loop_gate = session.scalar( + select(TI) + .where(TI.id == task_instance_id, TI.working_set.is_(True)) + .with_for_update(of=TI) + .execution_options(populate_existing=True) + ) + if loop_gate is not None: + try: + loop_context = TaskCoordinateResolver(dag_bag, session).loop_context(loop_gate) + except TaskNotFound: + # A gate renamed mid-run on an unversioned bundle leaves this attempt without a loop. + log.warning("Loop gate is missing from its pinned Dag version; failing it") + if isinstance(ti_patch_payload, TISuccessStatePayload): + ti_patch_payload = _build_rejected_gate_payload( + ti_patch_payload, + reason="Loop gate is missing from its pinned Dag version", + retry=False, + ) + loop_context = None + if loop_context is None: + loop_gate = None + elif loop_context[0].gate_task_id != loop_gate.task_id: + raise HTTPException(status_code=409, detail={"reason": "invalid_loop_gate"}) + else: + loop_group = loop_context[0] + old = ( select( TI.state, @@ -522,6 +578,25 @@ def ti_update_state( detail={"reason": "invalid_partition_key", "message": str(e)}, ) from e + gate_completed = False + if ( + loop_gate is not None + and loop_group is not None + and isinstance(ti_patch_payload, TISuccessStatePayload) + ): + try: + loop_gate.dag_run.complete_loop_gate( + loop_gate, loop_group, TaskInstanceState.SUCCESS, session=session + ) + gate_completed = True + except InvalidLoopDecision as error: + log.warning("Loop gate success rejected", error=str(error)) + ti_patch_payload = _build_rejected_gate_payload( + ti_patch_payload, + reason=f"Loop gate success rejected: {error}", + retry=_is_eligible_to_retry(previous_state, try_number, max_tries), + ) + # We exclude_unset to avoid updating fields that are not set in the payload data = ti_patch_payload.model_dump( exclude={"task_outlets", "outlet_events", "retry_delay_seconds", "retry_reason"}, @@ -545,6 +620,8 @@ def ti_update_state( # Let DataErrorHandler return a 422 instead of silently marking the TI FAILED below. raise except Exception: + if loop_gate is not None: + raise # Set a task to failed in case any unexpected exception happened during task state update log.exception( "Error updating Task Instance state. Setting the task to failed.", @@ -595,6 +672,9 @@ def ti_update_state( # Defer to app-level SQLAlchemyError handler (returns HTTP 500). raise + if loop_gate is not None and loop_group is not None and not gate_completed: + loop_gate.dag_run.complete_loop_gate(loop_gate, loop_group, updated_state, session=session) + if updated_state == TaskInstanceState.SUCCESS: if conf.getboolean("state_store", "clear_on_success"): scope = TaskScope( @@ -628,6 +708,15 @@ def ti_update_state( callback() +def _build_rejected_gate_payload( + payload: TISuccessStatePayload, *, reason: str, retry: bool +) -> TIRetryStatePayload | TITerminalStatePayload: + carried = payload.model_dump(include={"end_date", "rendered_map_index"}, exclude_unset=True) + if retry: + return TIRetryStatePayload(state=IntermediateTIState.UP_FOR_RETRY, retry_reason=reason, **carried) + return TITerminalStatePayload(state=TerminalStateNonSuccess.FAILED, retry_reason=reason, **carried) + + def _emit_task_span(ti, state, *, resolver: TaskCoordinateResolver): # just to be safe if not ti.dag_run: @@ -941,29 +1030,27 @@ def ti_skip_downstream( for task_id, index in targets if not resolver.has_regions(dag_id, run_id, task_id) ] - selected_ids = { - ti.id + selects = [ + resolver.select_skip_target_ids(caller=caller, task_id=task_id, map_indexes=index) for task_id, index in targets if (task_id, index) not in plain - for ti in resolver.resolve( - dag_id=dag_id, run_id=run_id, task_id=task_id, caller=caller, map_indexes=index - ) - } + ] if plain: - selected_ids.update( - session.scalars( - resolver.select_legacy_task_ids( - dag_id=dag_id, - run_id=run_id, - task_ids=[task_id for task_id, index in plain if index is None], - slots=[(task_id, index) for task_id, index in plain if index is not None], - ) + selects.append( + resolver.select_legacy_task_ids( + dag_id=dag_id, + run_id=run_id, + task_ids=[task_id for task_id, index in plain if index is None], + slots=[(task_id, index) for task_id, index in plain if index is not None], ) ) except AmbiguousProducerError as error: raise HTTPException(status.HTTP_409_CONFLICT, str(error)) from error except ValueError as error: raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error + if not selects: + return + selected_ids = set(session.scalars(union(*selects))) # Don't overwrite tasks that are already executing or finished. # See: https://github.com/apache/airflow/issues/59378 @@ -981,6 +1068,7 @@ def ti_skip_downstream( query = ( update(TI) .where( + TI.working_set.is_(True), TI.id.in_(selected_ids), skippable_state_clause, ) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py index c8d96e5c5d777..02fb6eceb6b11 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py @@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Annotated from uuid import UUID +import attrs from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, status from pydantic import JsonValue from sqlalchemy import select @@ -40,8 +41,14 @@ from airflow.api_fastapi.execution_api.security import CurrentTIToken from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyArchivedTaskStateUpdates from airflow.exceptions import TaskNotFound -from airflow.models.dagrun import DagRun -from airflow.models.dynamic_region import SENTINEL_REGION_ID, AmbiguousProducerError +from airflow.models.dagrun import DagRun, InvalidLoopDecision, validate_loop_decision +from airflow.models.dynamic_region import ( + LOOP_DECISION_KEY, + LOOP_XCOM_PREFIX, + SENTINEL_REGION_ID, + AmbiguousProducerError, + build_loop_sequence_order, +) from airflow.models.task_coordinates import ( TaskCoordinateResolver, build_coordinate_filters, @@ -50,9 +57,12 @@ from airflow.models.taskinstance import TaskInstance from airflow.models.xcom import XCOM_RETURN_KEY, XComModel, xcom_entity from airflow.utils.db import get_query_count +from airflow.utils.state import TaskInstanceState if TYPE_CHECKING: from sqlalchemy.orm import Session + from sqlalchemy.orm.util import AliasedClass + from sqlalchemy.sql.elements import ColumnElement from airflow.api_fastapi.execution_api.datamodels.token import TIToken from airflow.models.dagbag import DBDagBag @@ -136,6 +146,23 @@ def has_xcom_access( log = logging.getLogger(__name__) +@attrs.frozen +class _XComRead: + """The XCom rows a read selects, and whether they are the executions of a loop task seen from outside it.""" + + statement: Select + loop_sequence: bool = False + + def get_order(self, entity: AliasedClass, *, descending: bool = False) -> tuple[ColumnElement, ...]: + """Order the rows by their place in the sequence the read returns, or reverse that order.""" + columns = ( + build_loop_sequence_order(entity.region_id, entity.region_index) + if self.loop_sequence + else (entity.map_index,) + ) + return tuple(column.desc() if descending else column.asc() for column in columns) + + def _build_xcom_read( *, dag_id: str, @@ -147,8 +174,15 @@ def _build_xcom_read( token: TIToken, map_index: int | None = None, include_prior_dates: bool = False, -) -> Select: - """Select the XCom rows of the producers visible to the calling task instance.""" + previous_iteration: bool = False, + sequence: bool = False, +) -> _XComRead: + """ + Select the XCom rows of the producers visible to the calling task instance. + + A ``sequence`` read returns several values in order. When it targets a loop task from outside that + loop, it covers every iteration of the task rather than the one a caller inside the loop would see. + """ resolver = TaskCoordinateResolver(dag_bag, session) read = partial( XComModel.get_many, @@ -159,8 +193,12 @@ def _build_xcom_read( include_prior_dates=include_prior_dates, ) try: - if not resolver.has_regions(dag_id, None if include_prior_dates else run_id, task_id): - return read(region_id=SENTINEL_REGION_ID, map_indexes=map_index) + if previous_iteration and include_prior_dates: + raise ValueError("previous_iteration cannot be combined with prior dates") + if not previous_iteration and not resolver.has_regions( + dag_id, None if include_prior_dates else run_id, task_id + ): + return _XComRead(read(region_id=SENTINEL_REGION_ID, map_indexes=map_index)) caller = session.get(TaskInstance, token.id) if include_prior_dates: task = None @@ -176,18 +214,27 @@ def _build_xcom_read( ) if task is not None and enclosing_loop(task) is not None: raise ValueError("Prior-date lookups are not supported for tasks inside a loop") - return read(region_id=SENTINEL_REGION_ID, include_node_regions=True, map_indexes=map_index) - producers = resolver.resolve( - dag_id=dag_id, - run_id=run_id, - task_id=task_id, - caller=caller, - map_indexes=map_index, + return _XComRead( + read(region_id=SENTINEL_REGION_ID, include_node_regions=True, map_indexes=map_index) + ) + loop_sequence = sequence and resolver.is_loop_task_read_from_outside( + dag_id=dag_id, run_id=run_id, task_id=task_id, caller=caller ) - return XComModel.get_many( - run_id=run_id, - key=key, - producer_ids=select(TaskInstance.id).where(TaskInstance.id.in_([ti.id for ti in producers])), + return _XComRead( + XComModel.get_many( + run_id=run_id, + key=key, + producer_ids=resolver.select_producer_ids( + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + caller=caller, + map_indexes=map_index, + previous_iteration=previous_iteration, + all_iterations=loop_sequence, + ), + ), + loop_sequence=loop_sequence, ) except AmbiguousProducerError as error: raise HTTPException(status.HTTP_409_CONFLICT, str(error)) from error @@ -203,6 +250,7 @@ def xcom_query( session: SessionDep, dag_bag: DagBagDep, map_index: Annotated[int | None, Query()] = None, + previous_iteration: bool = False, token=CurrentTIToken, ) -> Select: return _build_xcom_read( @@ -214,7 +262,9 @@ def xcom_query( dag_bag=dag_bag, token=token, map_index=map_index, - ) + previous_iteration=previous_iteration, + sequence=True, + ).statement @router.get( @@ -235,6 +285,7 @@ def get_mapped_xcom_by_index( offset: int, session: SessionDep, dag_bag: DagBagDep, + previous_iteration: bool = False, token=CurrentTIToken, ) -> XComSequenceIndexResponse: xcom_read = _build_xcom_read( @@ -245,14 +296,15 @@ def get_mapped_xcom_by_index( session=session, dag_bag=dag_bag, token=token, + previous_iteration=previous_iteration, + sequence=True, ) - entity = xcom_entity(xcom_read) - xcom_query = xcom_read - xcom_query = xcom_query.order_by(None) + entity = xcom_entity(xcom_read.statement) + xcom_query = xcom_read.statement.order_by(None) if offset >= 0: - xcom_query = xcom_query.order_by(entity.map_index.asc()).offset(offset) + xcom_query = xcom_query.order_by(*xcom_read.get_order(entity)).offset(offset) else: - xcom_query = xcom_query.order_by(entity.map_index.desc()).offset(-1 - offset) + xcom_query = xcom_query.order_by(*xcom_read.get_order(entity, descending=True)).offset(-1 - offset) result: tuple[XComModel] | None if (result := session.scalars(xcom_query.limit(1)).first()) is None: @@ -273,6 +325,7 @@ class GetXComSliceFilterParams(BaseModel): stop: int | None = None step: int | None = None include_prior_dates: bool = False + previous_iteration: bool = False def _get_sliced_query_or_empty(query: Select, low: int, high: int) -> Select: @@ -310,10 +363,12 @@ def get_mapped_xcom_by_slice( dag_bag=dag_bag, token=token, include_prior_dates=params.include_prior_dates, + previous_iteration=params.previous_iteration, + sequence=True, ) - entity = xcom_entity(xcom_read) - query = xcom_read - query = query.order_by(None) + entity = xcom_entity(xcom_read.statement) + query = xcom_read.statement.order_by(None) + ascending, descending = xcom_read.get_order(entity), xcom_read.get_order(entity, descending=True) step = params.step or 1 @@ -323,25 +378,25 @@ def get_mapped_xcom_by_slice( if (start := params.start) is None: if (stop := params.stop) is None: if step >= 0: - query = query.order_by(entity.map_index.asc()) + query = query.order_by(*ascending) else: - query = query.order_by(entity.map_index.desc()) + query = query.order_by(*descending) step = -step elif stop >= 0: - query = query.order_by(entity.map_index.asc()) + query = query.order_by(*ascending) if step >= 0: query = query.limit(stop) else: query = query.offset(stop + 1) else: - query = query.order_by(entity.map_index.desc()) + query = query.order_by(*descending) step = -step if step > 0: query = query.limit(-stop - 1) else: query = query.offset(-stop) elif start >= 0: - query = query.order_by(entity.map_index.asc()) + query = query.order_by(*ascending) if (stop := params.stop) is None: if step >= 0: query = query.offset(start) @@ -355,7 +410,7 @@ def get_mapped_xcom_by_slice( else: query = _get_sliced_query_or_empty(query, stop + 1, start + 1) else: - query = query.order_by(entity.map_index.desc()) + query = query.order_by(*descending) step = -step if (stop := params.stop) is None: if step > 0: @@ -419,6 +474,7 @@ class GetXcomFilterParams(BaseModel): map_index: int = -1 include_prior_dates: bool = False offset: int | None = None + previous_iteration: bool = False @router.get( @@ -482,15 +538,19 @@ def get_xcom( token=token, map_index=params.map_index if params.offset is None else None, include_prior_dates=params.include_prior_dates, + previous_iteration=params.previous_iteration, + sequence=params.offset is not None, ) - entity = xcom_entity(xcom_read) - xcom_query = xcom_read + entity = xcom_entity(xcom_read.statement) + xcom_query = xcom_read.statement if params.offset is not None: xcom_query = xcom_query.where(entity.value.is_not(None)).order_by(None) if params.offset >= 0: - xcom_query = xcom_query.order_by(entity.map_index.asc()).offset(params.offset) + xcom_query = xcom_query.order_by(*xcom_read.get_order(entity)).offset(params.offset) else: - xcom_query = xcom_query.order_by(entity.map_index.desc()).offset(-1 - params.offset) + xcom_query = xcom_query.order_by(*xcom_read.get_order(entity, descending=True)).offset( + -1 - params.offset + ) # We use `BaseXCom.get_many` to fetch XComs directly from the database, bypassing the XCom Backend. # This avoids deserialization via the backend (e.g., from a remote storage like S3) and instead @@ -528,7 +588,10 @@ def get_xcom( status.HTTP_400_BAD_REQUEST, "The key is empty, the value is too large to map, or is unserializable", ), - (status.HTTP_409_CONFLICT, "Several live task instances match the producer"), + ( + status.HTTP_409_CONFLICT, + "Several live task instances match the producer, or the loop gate is no longer running", + ), ] ), ) @@ -538,6 +601,7 @@ def set_xcom( task_id: str, key: Annotated[str, Path(min_length=1)], session: SessionDep, + dag_bag: DagBagDep, value: Annotated[ JsonValue, Body( @@ -568,6 +632,45 @@ def set_xcom( """Set an Airflow XCom.""" from airflow.configuration import conf + region_id: UUID | None = None + region_index: int | None = None + if key == LOOP_DECISION_KEY: + caller = session.get(TaskInstance, token.id) + if ( + value not in ("continue", "stop") + or mapped_length is not None + or dag_result + or map_index != -1 + or caller is None + or (dag_id, run_id, task_id) != (caller.dag_id, caller.run_id, caller.task_id) + ): + raise HTTPException(status.HTTP_400_BAD_REQUEST, "Invalid loop decision write") + session.execute( + select(DagRun) + .where(DagRun.dag_id == caller.dag_id, DagRun.run_id == caller.run_id) + .with_for_update() + ).scalar_one() + caller = session.scalar( + select(TaskInstance) + .where(TaskInstance.id == token.id) + .with_for_update() + .execution_options(populate_existing=True) + ) + if caller is None: + raise HTTPException(status.HTTP_409_CONFLICT, "Loop gate is no longer running") + loop = TaskCoordinateResolver(dag_bag, session).loop_context(caller) + if loop is None or loop[0].gate_task_id != caller.task_id: + raise HTTPException(status.HTTP_400_BAD_REQUEST, "Loop decisions require the pinned loop gate") + if caller.state != TaskInstanceState.RUNNING or caller.working_set is not True: + raise HTTPException(status.HTTP_409_CONFLICT, "Loop gate is no longer running") + try: + validate_loop_decision(value, caller.region_index, loop[0]) + except InvalidLoopDecision as error: + raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error + region_id, region_index = caller.region_id, caller.region_index + elif key.startswith(LOOP_XCOM_PREFIX): + raise HTTPException(status.HTTP_400_BAD_REQUEST, "Reserved loop XCom key") + # Validate that the provided key is not empty # XCom keys must be non-empty strings to ensure proper data retrieval and avoid ambiguity. if not key: @@ -614,6 +717,8 @@ def set_xcom( run_id=run_id, task_id=task_id, map_index=map_index, + region_id=region_id, + region_index=region_index, session=session, ), serialize=False, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index b87c99b7c27a1..f23972db0045d 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -56,12 +56,16 @@ AddArgBindingsToTIRunContext, AddCallbackRunEndpoint, AddDagRunNoteUpdateEndpoint, + AddLoopContext, AddMultiTeamToTIRunContext, AddStoppedTaskReport, AddTaskInstanceRegionCoordinates, AddTerminalStateRetryReasonField, IdentifyArchivedTaskStateUpdates, ) +from airflow.api_fastapi.execution_api.versions.v2026_10_30_xcom_params import ( + AddPreviousIterationToXComFilterParams, +) bundle = VersionBundle( HeadVersion(), @@ -70,8 +74,10 @@ AddArgBindingsToTIRunContext, AddCallbackRunEndpoint, AddDagRunNoteUpdateEndpoint, + AddLoopContext, AddTerminalStateRetryReasonField, AddMultiTeamToTIRunContext, + AddPreviousIterationToXComFilterParams, AddStoppedTaskReport, AddTaskInstanceRegionCoordinates, IdentifyArchivedTaskStateUpdates, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py index fd8ca698ed3b7..a75c7591982b2 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py @@ -36,6 +36,17 @@ ) +class AddLoopContext(VersionChange): + """Provide the pinned loop definition and enclosing iteration at task startup.""" + + description = __doc__ + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("loop").didnt_exist,) + + @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] + def remove_loop_context(response: ResponseInfo) -> None: # type: ignore[misc] + response.body.pop("loop", None) + + class AddTaskInstanceRegionCoordinates(VersionChange): """Carry task execution coordinates independently of public mapped position, omitting them outside a region.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30_xcom_params.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30_xcom_params.py new file mode 100644 index 0000000000000..4a97ff9af91e7 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30_xcom_params.py @@ -0,0 +1,35 @@ +# 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. + +from __future__ import annotations + +from cadwyn import VersionChange, schema + +# These models live next to the routes that use them, and the routes import this version package, so the +# changes that touch them cannot sit in ``v2026_10_30`` without a circular import. +from airflow.api_fastapi.execution_api.routes.xcoms import GetXcomFilterParams, GetXComSliceFilterParams + + +class AddPreviousIterationToXComFilterParams(VersionChange): + """Add the `previous_iteration` field to GetXComSliceFilterParams and GetXcomFilterParams.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(GetXComSliceFilterParams).field("previous_iteration").didnt_exist, + schema(GetXcomFilterParams).field("previous_iteration").didnt_exist, + ) diff --git a/airflow-core/docs/authoring-and-scheduling/examples/example_task_loops.py b/airflow-core/src/airflow/example_dags/example_task_loops.py similarity index 100% rename from airflow-core/docs/authoring-and-scheduling/examples/example_task_loops.py rename to airflow-core/src/airflow/example_dags/example_task_loops.py diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index ca0604af6b602..5c5857ca8ecad 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -114,7 +114,7 @@ from airflow.models.log import resolve_team_name from airflow.models.pool import normalize_pool_name_for_stats from airflow.models.serialized_dag import SerializedDagModel -from airflow.models.task_coordinates import TaskCoordinateResolver +from airflow.models.task_coordinates import LOOP_GATE_OPERATOR, TaskCoordinateResolver from airflow.models.taskinstance import TaskInstance from airflow.models.team import Team from airflow.models.trigger import TRIGGER_FAIL_REPR, Trigger, TriggerFailureReason, handle_event_submit @@ -1685,6 +1685,11 @@ def process_executor_events( task = dag.get_task(ti.task_id) except Exception: + if state == TaskInstanceState.SUCCESS and ti.operator == LOOP_GATE_OPERATOR: + cls.logger().exception("Failing loop gate %s: %s", ti, msg) + ti.task = None + ti.handle_failure(error=msg, session=session) + continue cls.logger().exception("Marking task instance %s as %s", ti, state) ti.set_state(state) continue diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index 8d4ebcc91d389..0643e613babed 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -88,7 +88,12 @@ from airflow.models.backfill import Backfill from airflow.models.base import Base, StringID from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel -from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion, public_region_filter +from airflow.models.dynamic_region import ( + LOOP_DECISION_KEY, + SENTINEL_REGION_ID, + DynamicRegion, + public_region_filter, +) from airflow.models.taskinstance import TaskInstance as TI, _add_and_prime_mapped_ti, clear_task_instances from airflow.models.tasklog import LogTemplate from airflow.serialization.definitions.deadline import SerializedReferenceModels @@ -127,6 +132,7 @@ from airflow.sdk import DAG as SDKDAG from airflow.serialization.definitions.dag import SerializedDAG from airflow.serialization.definitions.mappedoperator import Operator + from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup from airflow.timetables.base import Timetable CreatedTasks = TypeVar("CreatedTasks", Iterator["dict[str, Any]"], Iterator[TI]) @@ -141,6 +147,21 @@ tracer = trace.get_tracer(__name__) +class InvalidLoopDecision(ValueError): + """A gate cannot complete successfully with the available decision.""" + + +def validate_loop_decision(decision: str | None, region_index: int, group: SerializedLoopTaskGroup) -> None: + """Raise :class:`InvalidLoopDecision` unless ``decision`` is valid for the gate of pass ``region_index``.""" + at_limit = region_index + 1 >= group.max_iterations + if decision not in ("continue", "stop"): + raise InvalidLoopDecision("Successful loop gate requires a decision") + if decision == "continue" and at_limit: + raise InvalidLoopDecision("Loop cannot continue beyond its iteration limit") + if not group.has_until and (decision == "stop") != at_limit: + raise InvalidLoopDecision("Fixed-count loop decision does not match its iteration limit") + + class TISchedulingDecision(NamedTuple): """Type of return for DagRun.task_instance_scheduling_decisions.""" @@ -1931,11 +1952,10 @@ def task_filter(task: Operator) -> bool: # Create the missing tasks, including mapped tasks missing_tasks = [task for task in dag.task_dict.values() if task_filter(task)] tis_to_create = list( - self._create_tasks( + self._create_initial_tasks( missing_tasks, task_creator, session=session, - expand_literals=True, ) ) self._create_task_instances(self.dag_id, tis_to_create, created_counts, hook_is_noop, session=session) @@ -2205,6 +2225,121 @@ def create_ti(task: Operator, indexes: Iterable[int], region_id: UUID) -> Iterat creator = create_ti return creator + def complete_loop_gate( + self, + gate: TI, + group: SerializedLoopTaskGroup, + state: TaskInstanceState, + *, + session: Session, + ) -> None: + """Consume a gate decision while the caller holds the DagRun and TI locks.""" + from airflow.models.xcom import XComModelV2 + from airflow.settings import task_instance_mutation_hook + + if gate.dag_version_id is None: + raise InvalidLoopDecision("Loop gate requires a pinned DAG version") + signal = XComModelV2.get_for_attempt(gate.id, LOOP_DECISION_KEY, session=session) + later_gate = session.scalar( + select(TI.id) + .where( + TI.working_set.is_(True), + TI.dag_id == self.dag_id, + TI.run_id == self.run_id, + TI.task_id == gate.task_id, + TI.region_id == gate.region_id, + TI.region_index > gate.region_index, + ) + .limit(1) + ) + if state != TaskInstanceState.SUCCESS or later_gate: + if signal is not None: + session.delete(signal) + return + decision = signal.value if signal is not None else None + validate_loop_decision(decision, gate.region_index, group) + session.delete(signal) + if decision == "stop": + return + created_counts: dict[str, int] = defaultdict(int) + hook_is_noop: Literal[True, False] = getattr(task_instance_mutation_hook, "is_noop", False) + creator = self._get_task_creator( + created_counts, task_instance_mutation_hook, hook_is_noop, gate.dag_version_id + ) + tasks = list( + self._create_tasks( + group.iter_tasks(), + creator, + session=session, + parent_region=(gate.region_id, gate.region_index + 1), + ) + ) + self._create_task_instances( + self.dag_id, tasks, created_counts, hook_is_noop, session=session, propagate_errors=True + ) + + def _create_initial_tasks( + self, + tasks: Iterable[Operator], + task_creator: Callable[[Operator, Iterable[int], UUID], CreatedTasks], + *, + session: Session, + ) -> CreatedTasks: + from airflow.models.task_coordinates import enclosing_loop + + ordinary_tasks = [] + loop_tasks = defaultdict(list) + loops = {} + for task in tasks: + if loop := enclosing_loop(task): + loops[loop.group_id] = loop + loop_tasks[loop.group_id].append(task) + else: + ordinary_tasks.append(task) + yield from self._create_tasks(ordinary_tasks, task_creator, session=session, expand_literals=True) + for group_id, members in loop_tasks.items(): + coordinates = ( + session.execute( + select(TI.region_id, TI.region_index).where( + TI.working_set.is_(True), + TI.dag_id == self.dag_id, + TI.run_id == self.run_id, + TI.task_id == loops[group_id].gate_task_id, + ) + ) + .tuples() + .all() + ) + if not coordinates: + if session.scalar( + select(DynamicRegion.id) + .where( + DynamicRegion.dag_id == self.dag_id, + DynamicRegion.run_id == self.run_id, + DynamicRegion.node_id == group_id, + ) + .limit(1) + ): + log.warning( + "skipping loop that has regions but no live gate", + dag_id=self.dag_id, + run_id=self.run_id, + loop=group_id, + ) + continue + region = DynamicRegion.get_or_create( + dag_id=self.dag_id, run_id=self.run_id, node_id=loops[group_id].node_id, session=session + ) + coordinates = [(region.id, 0)] + for region_id, region_index in coordinates: + yield from self._create_tasks( + members, + task_creator, + session=session, + parent_region=(region_id, region_index), + expand_literals=True, + ) + def _create_tasks( self, tasks: Iterable[Operator], @@ -2251,6 +2386,7 @@ def _create_task_instances( hook_is_noop: bool, *, session: Session, + propagate_errors: bool = False, ) -> None: """ Create the necessary task instances from the given tasks. @@ -2280,6 +2416,8 @@ def _create_task_instances( ) session.flush() except (IntegrityError, StaleDataError) as exc: + if propagate_errors: + raise self.log.info( "Hit %s while creating the TIs for %s- %s", type(exc).__name__, diff --git a/airflow-core/src/airflow/models/dynamic_region.py b/airflow-core/src/airflow/models/dynamic_region.py index e9073047bf3e8..fbef2054312bd 100644 --- a/airflow-core/src/airflow/models/dynamic_region.py +++ b/airflow-core/src/airflow/models/dynamic_region.py @@ -20,7 +20,8 @@ import struct from collections.abc import Collection, Iterable from datetime import datetime -from typing import TYPE_CHECKING +from functools import partial +from typing import TYPE_CHECKING, Any from uuid import UUID import attrs @@ -35,6 +36,8 @@ UniqueConstraint, and_, event, + false, + func, literal, or_, select, @@ -48,8 +51,11 @@ from airflow.utils.sqlalchemy import CompactUUID, UtcDateTime SENTINEL_REGION_ID = UUID(int=0) +LOOP_XCOM_PREFIX = "_airflow_loop_" +LOOP_DECISION_KEY = f"{LOOP_XCOM_PREFIX}decision" if TYPE_CHECKING: + from sqlalchemy import Select from sqlalchemy.engine import Connection from sqlalchemy.orm import Mapper, Session from sqlalchemy.sql.elements import ColumnElement @@ -466,29 +472,88 @@ def _match_any(column, value: str | Collection[str]) -> ColumnElement[bool]: return column == value if isinstance(value, str) else column.in_(value) -def resolve_current_producers( +def _filter_producers( + query: Select, *, dag_id: str, run_id: str, task_id: str, is_mapped: bool, - context: ProducerContext | None = None, - map_indexes: int | Collection[int] | None = None, - region_id: UUID | None = None, - region_index: int | None = None, - session: Session, -) -> tuple[TaskInstance, ...]: - """Resolve the live producer task instances whose data the caller reads by task instance UUID.""" + map_indexes: int | Collection[int] | None, + region_id: UUID | None, + region_index: int | None, + top_level_only: bool, + loop_pass: _LoopPassRegions | None = None, +) -> Select: + """Push every coordinate predicate into SQL.""" from airflow.models.taskinstance import TaskInstance + query = query.where( + TaskInstance.dag_id == dag_id, + TaskInstance.run_id == run_id, + TaskInstance.task_id == task_id, + TaskInstance.working_set.is_(True), + ) + if region_id is not None: + query = query.where(TaskInstance.region_id == region_id) + if region_index is not None: + query = query.where(TaskInstance.region_index == region_index) + if loop_pass is not None: + query = query.where(_build_loop_pass_filter(loop_pass)) + if top_level_only: + if is_mapped: + top_level_region = ( + select(DynamicRegion.id) + .where(DynamicRegion.id == TaskInstance.region_id, DynamicRegion.parent_region_id.is_(None)) + .correlate(TaskInstance) + .exists() + ) + query = query.where(or_(TaskInstance.region_id == SENTINEL_REGION_ID, top_level_region)) + else: + query = query.where(TaskInstance.region_id == SENTINEL_REGION_ID, TaskInstance.region_index == -1) + if map_indexes is None: + return query + if not is_mapped: + wanted = map_indexes == -1 if isinstance(map_indexes, int) else -1 in map_indexes + return query if wanted else query.where(false()) + if isinstance(map_indexes, int): + return query.where(TaskInstance.region_index == map_indexes) + if isinstance(map_indexes, range) and map_indexes.step == 1: + return query.where( + TaskInstance.region_index >= map_indexes.start, TaskInstance.region_index < map_indexes.stop + ) + return query.where(TaskInstance.region_index.in_(list(map_indexes))) + + +def _validate_producer_request( + *, context: ProducerContext | None, region_id: UUID | None, region_index: int | None +) -> None: if region_index is not None and region_id is None: raise ValueError("region_index requires an explicit producer region_id") if context and context.previous_iteration and context.loop_node_id is None: raise ValueError("Previous-iteration lookup requires a loop context") + + +def _build_candidate_query( + query: Select[Any], + *, + dag_id: str, + run_id: str, + task_id: str, + is_mapped: bool, + context: ProducerContext | None, + map_indexes: int | Collection[int] | None, + region_id: UUID | None, + region_index: int | None, + session: Session, +) -> Select[Any] | None: + """Narrow ``query`` to the candidate producers of the caller's scope, or ``None`` if there are none.""" regions: dict[UUID, DynamicRegion] = {} position = None if context and region_id is None: - regions = load_region_ancestry({context.region_id}, dag_id=dag_id, run_id=run_id, session=session) + regions = load_region_ancestry( + {context.region_id} - {SENTINEL_REGION_ID}, dag_id=dag_id, run_id=run_id, session=session + ) if context.loop_node_id is not None: position = loop_position(regions, context.region_id, context.region_index, context.loop_node_id) if position is None: @@ -496,47 +561,239 @@ def resolve_current_producers( if context.previous_iteration: position = position[0], position[1] - 1 if position[1] < 0: - return () - query = select(TaskInstance).where( - TaskInstance.dag_id == dag_id, - TaskInstance.run_id == run_id, - TaskInstance.task_id == task_id, - TaskInstance.working_set.is_(True), + return None + + return _filter_producers( + query, + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + is_mapped=is_mapped, + map_indexes=map_indexes, + region_id=region_id, + region_index=region_index, + top_level_only=region_id is None and position is None, + loop_pass=_load_loop_pass_regions(position, session=session) if position is not None else None, ) - if region_id is not None: - query = query.where(TaskInstance.region_id == region_id) - if region_index is not None: - query = query.where(TaskInstance.region_index == region_index) - if position is not None: - query = query.where(_build_loop_pass_filter(_load_loop_pass_regions(position, session=session))) + + +def _check_producers_unambiguous( + region_indexes: Iterable[int], *, dag_id: str, run_id: str, task_id: str, is_mapped: bool +) -> None: + seen: set[int] = set() + for region_index in region_indexes: + public_index = region_index if is_mapped else -1 + if public_index in seen: + raise AmbiguousProducerError( + f"Multiple live producers for {dag_id}/{run_id}/{task_id} index {public_index}" + ) + seen.add(public_index) + + +def resolve_current_producers( + *, + dag_id: str, + run_id: str, + task_id: str, + is_mapped: bool, + context: ProducerContext | None = None, + map_indexes: int | Collection[int] | None = None, + region_id: UUID | None = None, + region_index: int | None = None, + session: Session, +) -> tuple[TaskInstance, ...]: + """Resolve the live producer task instances whose data the caller reads by task instance UUID.""" + from airflow.models.taskinstance import TaskInstance + + _validate_producer_request(context=context, region_id=region_id, region_index=region_index) + query = _build_candidate_query( + select(TaskInstance), + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + is_mapped=is_mapped, + context=context, + map_indexes=map_indexes, + region_id=region_id, + region_index=region_index, + session=session, + ) + if query is None: + return () candidates = session.scalars(query).all() - if region_id is None and position is None: - regions.update( - load_region_ancestry( - {ti.region_id for ti in candidates} - set(regions), + _check_producers_unambiguous( + (ti.region_index for ti in candidates), + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + is_mapped=is_mapped, + ) + return tuple(sorted(candidates, key=lambda ti: ti.region_index if is_mapped else -1)) + + +def select_current_producer_ids( + *, + dag_id: str, + run_id: str, + task_id: str, + is_mapped: bool, + context: ProducerContext | None = None, + map_indexes: int | Collection[int] | None = None, + region_id: UUID | None = None, + region_index: int | None = None, + session: Session, +) -> Select[tuple[UUID]]: + """ + Select the ids of the live producers :func:`resolve_current_producers` would return. + + A task whose live executions all share one region cannot have two producers at one index, so + that common case stays a pure SQL selection whose cost does not depend on the number of + mapped instances. Everything else resolves the rows and pins their ids. + """ + from airflow.models.taskinstance import TaskInstance + + _validate_producer_request(context=context, region_id=region_id, region_index=region_index) + filter_producers = partial( + _filter_producers, + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + is_mapped=is_mapped, + map_indexes=map_indexes, + region_id=region_id, + region_index=region_index, + top_level_only=region_id is None, + ) + query = filter_producers(select(TaskInstance.id)) + if region_id is None: + if context is not None: + single_region = False + else: + first_region = filter_producers(select(TaskInstance.region_id)).limit(1) + other_region = filter_producers(select(TaskInstance.id)).where( + TaskInstance.region_id != first_region.scalar_subquery() + ) + single_region = not session.scalar(select(other_region.exists())) + if not single_region: + candidate_query = _build_candidate_query( + select(TaskInstance.id, TaskInstance.region_index), dag_id=dag_id, run_id=run_id, + task_id=task_id, + is_mapped=is_mapped, + context=context, + map_indexes=map_indexes, + region_id=None, + region_index=None, session=session, ) + candidates = session.execute(candidate_query).all() if candidate_query is not None else [] + _check_producers_unambiguous( + (candidate.region_index for candidate in candidates), + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + is_mapped=is_mapped, + ) + return select(TaskInstance.id).where( + TaskInstance.id.in_([candidate.id for candidate in candidates]) + ) + return query + + +def _load_loop_pass_owners( + *, dag_id: str, run_id: str, loop_node_id: str, session: Session +) -> list[tuple[UUID, int | None]]: + """ + Return each region that holds passes of a loop with the first pass it no longer owns. + + A fork takes over every pass from its ``resumes_from_index``, so a region owns only the passes + before the earliest fork that follows it. The newest region owns every pass from its own start. + """ + regions = session.scalars( + select(DynamicRegion).where( + DynamicRegion.dag_id == dag_id, + DynamicRegion.run_id == run_id, + DynamicRegion.node_id == loop_node_id, + ) + ).all() + originals = [region for region in regions if region.forked_from_region_id is None] + if len(originals) > 1: + raise ValueError(f"Loop {loop_node_id!r} has several executions in Dag run {run_id!r}") + successors = {region.forked_from_region_id: region for region in regions} + chain = originals[:] + while chain and chain[-1].id in successors: + chain.append(successors[chain[-1].id]) + owners: list[tuple[UUID, int | None]] = [] + cut: int | None = None + for region in reversed(chain): + owners.append((region.id, cut)) + cut = region.resumes_from_index if cut is None else min(cut, region.resumes_from_index) + return owners + + +def _build_pass_ownership_filter( + region_column, index_column, owners: Iterable[tuple[UUID, int | None]] +) -> ColumnElement[bool]: + return or_( + *( + region_column == region_id + if cut is None + else and_(region_column == region_id, index_column < cut) + for region_id, cut in owners ) + ) - selected: dict[int, TaskInstance] = {} - for ti in candidates: - if region_id is None and position is None: - if ti.region_id != SENTINEL_REGION_ID: - if not is_mapped or regions[ti.region_id].parent_region_id is not None: - continue - elif not is_mapped and ti.region_index != -1: - continue - public_index = ti.region_index if is_mapped else -1 - if isinstance(map_indexes, int): - if public_index != map_indexes: - continue - elif map_indexes is not None and public_index not in map_indexes: - continue - if public_index in selected: - raise AmbiguousProducerError( - f"Multiple live producers for {dag_id}/{run_id}/{task_id} index {public_index}" - ) - selected[public_index] = ti - return tuple(selected[index] for index in sorted(selected)) + +def select_loop_producer_ids( + *, dag_id: str, run_id: str, loop_node_id: str, task_id: str, is_mapped: bool, session: Session +) -> Select[tuple[UUID]]: + """ + Select every live execution of ``task_id`` across all passes of the loop ``loop_node_id``. + + That is what a task outside the loop sees of a loop task. A pass that a fork replaced does not count + even while its task instances are still live, because the fork's task instances stand in for it. + Order the executions with :func:`build_loop_sequence_order`. + """ + from airflow.models.taskinstance import TaskInstance + + query = select(TaskInstance.id).where( + TaskInstance.dag_id == dag_id, + TaskInstance.run_id == run_id, + TaskInstance.task_id == task_id, + TaskInstance.working_set.is_(True), + ) + owners = _load_loop_pass_owners(dag_id=dag_id, run_id=run_id, loop_node_id=loop_node_id, session=session) + if not owners: + return query.where(false()) + if not is_mapped: + return query.where( + _build_pass_ownership_filter(TaskInstance.region_id, TaskInstance.region_index, owners) + ) + pass_region = ( + select(DynamicRegion.id) + .where( + DynamicRegion.id == TaskInstance.region_id, + _build_pass_ownership_filter( + DynamicRegion.parent_region_id, DynamicRegion.parent_region_index, owners + ), + ) + .correlate(TaskInstance) + .exists() + ) + return query.where(pass_region) + + +def build_loop_sequence_order(region_id_column, region_index_column) -> tuple[ColumnElement[int], ...]: + """ + Order the rows :func:`select_loop_producer_ids` selects depth-first: by pass, then by map index. + + An unmapped task sits in the loop's own region, whose ``region_index`` is the pass. A mapped task sits + in a region nested under the pass, whose parent index is the pass and whose rows carry the map index. + """ + nested_pass = ( + select(DynamicRegion.parent_region_index) + .where(DynamicRegion.id == region_id_column) + .scalar_subquery() + ) + return func.coalesce(nested_pass, region_index_column), region_index_column diff --git a/airflow-core/src/airflow/models/task_coordinates.py b/airflow-core/src/airflow/models/task_coordinates.py index 7f7930192f1ff..40bdcc890792e 100644 --- a/airflow-core/src/airflow/models/task_coordinates.py +++ b/airflow-core/src/airflow/models/task_coordinates.py @@ -34,6 +34,8 @@ load_region_ancestry, loop_position, resolve_current_producers, + select_current_producer_ids, + select_loop_producer_ids, ) from airflow.models.taskinstance import TaskInstance from airflow.serialization.definitions.mappedoperator import is_mapped @@ -49,6 +51,23 @@ from airflow.serialization.definitions.dag import SerializedDAG, SerializedOperator +LOOP_GATE_OPERATOR = "LoopGateOperator" +"""Operator class name the Task SDK gives the gate task of a loop.""" + + +@attrs.frozen(kw_only=True) +class _ProducerRequest: + dag_id: str + run_id: str + task_id: str + is_mapped: bool + context: ProducerContext | None + map_indexes: int | Collection[int] | None + region_id: UUID | None + region_index: int | None + loop_node_id: str | None = None + + class TaskCoordinate(Protocol): """Stored task identity shared by live and retained task data.""" @@ -57,6 +76,7 @@ class TaskCoordinate(Protocol): task_id: str region_id: UUID region_index: int + dag_version_id: UUID | None def enclosing_loop(task: SerializedOperator) -> SerializedLoopTaskGroup | None: @@ -156,6 +176,21 @@ def adopt_dag(self, dag: SerializedDAG | None) -> None: if dag is not None and dag.dag_version_id is not None: self._dags.setdefault(dag.dag_version_id, dag) + def loop_context(self, ti: TaskCoordinate) -> tuple[SerializedLoopTaskGroup, int] | None: + if ti.region_id == SENTINEL_REGION_ID: + return None + task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id) + group = enclosing_loop(task) + if group is None: + return None + regions = load_region_ancestry( + [ti.region_id], dag_id=ti.dag_id, run_id=ti.run_id, session=self.session + ) + position = loop_position(regions, ti.region_id, ti.region_index, group.node_id) + if position is None: + raise ValueError("Task coordinates do not belong to the pinned loop") + return group, position[1] + def get_task( self, dag_id: str, run_id: str, task_id: str, *, dag_version_id: UUID | None = None ) -> SerializedOperator: @@ -371,39 +406,137 @@ def has_regions(self, dag_id: str, run_id: str | None, task_id: str) -> bool: query = query.where(TaskInstance.run_id == run_id) return bool(self.session.scalar(select(query.exists()))) - def resolve( + def resolve_dependency(self, caller: TaskInstance, task_id: str) -> tuple[TaskInstance, ...]: + try: + producer = self.get_task( + caller.dag_id, caller.run_id, task_id, dag_version_id=caller.dag_version_id + ) + except TaskNotFound: + return self.resolve(dag_id=caller.dag_id, run_id=caller.run_id, task_id=task_id, caller=caller) + loop = enclosing_loop(producer) + caller_task = self.get_task( + caller.dag_id, caller.run_id, caller.task_id, dag_version_id=caller.dag_version_id + ) + caller_loop = enclosing_loop(caller_task) + if ( + loop is not None + and loop.gate_task_id == task_id + and (caller_loop is None or caller_loop.group_id != loop.group_id) + ): + gates = self.session.scalars( + select(TaskInstance) + .join(DynamicRegion, DynamicRegion.id == TaskInstance.region_id) + .where( + TaskInstance.working_set.is_(True), + TaskInstance.dag_id == caller.dag_id, + TaskInstance.run_id == caller.run_id, + TaskInstance.task_id == task_id, + DynamicRegion.node_id == loop.group_id, + ) + .order_by(TaskInstance.region_index.desc()) + .limit(2) + ).all() + if len(gates) == 2 and gates[0].region_index == gates[1].region_index: + raise AmbiguousProducerError(f"Multiple live loop gates at pass {gates[0].region_index}") + return tuple(gates[:1]) + return self.resolve( + dag_id=caller.dag_id, run_id=caller.run_id, task_id=task_id, caller=caller, all_iterations=True + ) + + @staticmethod + def _filter_region_indexes( + query: Select, region_index: int | None, map_indexes: int | Collection[int] | None + ) -> Select: + if region_index is not None: + query = query.where(TaskInstance.region_index == region_index) + if isinstance(map_indexes, int): + query = query.where(TaskInstance.region_index == map_indexes) + elif map_indexes is not None: + query = query.where(TaskInstance.region_index.in_(map_indexes)) + return query + + def _filter_legacy_producers( self, + query: Select, *, dag_id: str, run_id: str, task_id: str, - caller: TaskInstance | None = None, - region_id: UUID | None = None, - region_index: int | None = None, - map_indexes: int | Collection[int] | None = None, - previous_iteration: bool = False, - ) -> tuple[TaskInstance, ...]: - if region_index is not None and region_id is None: - raise ValueError("region_index requires an explicit producer region_id") - if not previous_iteration and ( + region_index: int | None, + map_indexes: int | Collection[int] | None, + ) -> Select: + query = query.where( + TaskInstance.working_set.is_(True), + TaskInstance.dag_id == dag_id, + TaskInstance.run_id == run_id, + TaskInstance.task_id == task_id, + TaskInstance.region_id == SENTINEL_REGION_ID, + ) + return self._filter_region_indexes(query, region_index, map_indexes) + + def _filter_removed_task_producers( + self, + query: Select, + *, + dag_id: str, + run_id: str, + task_id: str, + region_id: UUID | None, + region_index: int | None, + map_indexes: int | Collection[int] | None, + ) -> Select: + """Find live rows of a task whose definition is gone, using only its own expansion regions.""" + query = query.join(DynamicRegion, DynamicRegion.id == TaskInstance.region_id).where( + TaskInstance.working_set.is_(True), + TaskInstance.dag_id == dag_id, + TaskInstance.run_id == run_id, + TaskInstance.task_id == task_id, + DynamicRegion.node_id == task_id, + ) + if region_id is not None: + query = query.where(TaskInstance.region_id == region_id) + return self._filter_region_indexes(query, region_index, map_indexes) + + def _is_legacy_lookup( + self, dag_id: str, run_id: str, task_id: str, region_id: UUID | None, previous_iteration: bool + ) -> bool: + return not previous_iteration and ( region_id == SENTINEL_REGION_ID or (region_id is None and not self.has_regions(dag_id, run_id, task_id)) - ): - query = select(TaskInstance).where( - TaskInstance.working_set.is_(True), - TaskInstance.dag_id == dag_id, - TaskInstance.run_id == run_id, - TaskInstance.task_id == task_id, - TaskInstance.region_id == SENTINEL_REGION_ID, - ) - if region_index is not None: - query = query.where(TaskInstance.region_index == region_index) - if isinstance(map_indexes, int): - query = query.where(TaskInstance.region_index == map_indexes) - elif map_indexes is not None: - query = query.where(TaskInstance.region_index.in_(map_indexes)) - return tuple(self.session.scalars(query.order_by(TaskInstance.region_index))) + ) + + def _is_outside_loop(self, loop: SerializedLoopTaskGroup, caller: TaskInstance) -> bool: + caller_loop = enclosing_loop( + self.get_task(caller.dag_id, caller.run_id, caller.task_id, dag_version_id=caller.dag_version_id) + ) + return caller_loop is None or caller_loop.group_id != loop.group_id + + def is_loop_task_read_from_outside( + self, *, dag_id: str, run_id: str, task_id: str, caller: TaskInstance | None + ) -> bool: + """Tell whether ``caller``, of the same Dag run and outside the loop, reads a task of that loop.""" + if caller is None or (caller.dag_id, caller.run_id) != (dag_id, run_id): + return False + try: + task = self.get_task(dag_id, run_id, task_id, dag_version_id=caller.dag_version_id) + except TaskNotFound: + return False + loop = enclosing_loop(task) + return loop is not None and self._is_outside_loop(loop, caller) + def _build_producer_request( + self, + *, + dag_id: str, + run_id: str, + task_id: str, + caller: TaskInstance | None, + region_id: UUID | None, + region_index: int | None, + map_indexes: int | Collection[int] | None, + previous_iteration: bool, + all_iterations: bool = False, + ) -> _ProducerRequest | None: shared_run = caller is not None and (caller.dag_id, caller.run_id) == (dag_id, run_id) try: task = ( @@ -414,30 +547,36 @@ def resolve( ) ) except TaskNotFound: - return self._resolve_removed_task( - dag_id=dag_id, - run_id=run_id, - task_id=task_id, - region_id=region_id, - region_index=region_index, - map_indexes=map_indexes, - ) + return None loop = enclosing_loop(task) context = None if region_id is None and loop is not None: if caller is None or (caller.dag_id, caller.run_id) != (dag_id, run_id): raise ValueError("A loop producer requires an explicit scope or a consumer inside the loop") - caller_loop = enclosing_loop( - self.get_task(dag_id, run_id, caller.task_id, dag_version_id=caller.dag_version_id) - ) - if caller_loop is None or caller_loop.group_id != loop.group_id: - raise ValueError("A loop producer requires an explicit scope or a consumer inside the loop") + if self._is_outside_loop(loop, caller): + if not all_iterations: + raise ValueError( + "A loop producer requires an explicit scope or a consumer inside the loop" + ) + if previous_iteration: + raise ValueError("Previous-iteration lookup requires a shared loop scope") + return _ProducerRequest( + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + is_mapped=task.get_needs_expansion(), + context=None, + map_indexes=map_indexes, + region_id=None, + region_index=None, + loop_node_id=loop.node_id, + ) context = ProducerContext( caller.region_id, caller.region_index, loop.group_id, previous_iteration ) if previous_iteration and context is None: raise ValueError("Previous-iteration lookup requires a shared loop scope") - return resolve_current_producers( + return _ProducerRequest( dag_id=dag_id, run_id=run_id, task_id=task_id, @@ -446,37 +585,187 @@ def resolve( map_indexes=map_indexes, region_id=region_id, region_index=region_index, - session=self.session, ) - def _resolve_removed_task( + def resolve( self, *, dag_id: str, run_id: str, task_id: str, - region_id: UUID | None, - region_index: int | None, - map_indexes: int | Collection[int] | None, + caller: TaskInstance | None = None, + region_id: UUID | None = None, + region_index: int | None = None, + map_indexes: int | Collection[int] | None = None, + previous_iteration: bool = False, + all_iterations: bool = False, ) -> tuple[TaskInstance, ...]: - """Find live rows of a task whose definition is gone, using only its own expansion regions.""" - query = ( - select(TaskInstance) - .join(DynamicRegion, DynamicRegion.id == TaskInstance.region_id) - .where( - TaskInstance.working_set.is_(True), - TaskInstance.dag_id == dag_id, - TaskInstance.run_id == run_id, - TaskInstance.task_id == task_id, - DynamicRegion.node_id == task_id, + """ + Resolve the live task instances of ``task_id`` that ``caller`` reads. + + With ``all_iterations``, a caller of the same Dag run that sits outside the loop of ``task_id`` + reads every live iteration of it instead of being refused. + """ + if region_index is not None and region_id is None: + raise ValueError("region_index requires an explicit producer region_id") + if self._is_legacy_lookup(dag_id, run_id, task_id, region_id, previous_iteration): + query = self._filter_legacy_producers( + select(TaskInstance), + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + region_index=region_index, + map_indexes=map_indexes, ) + return tuple(self.session.scalars(query.order_by(TaskInstance.region_index))) + request = self._build_producer_request( + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + caller=caller, + region_id=region_id, + region_index=region_index, + map_indexes=map_indexes, + previous_iteration=previous_iteration, + all_iterations=all_iterations, + ) + if request is not None and request.loop_node_id is not None: + return tuple( + self.session.scalars( + select(TaskInstance).where(TaskInstance.id.in_(self._select_loop_producer_ids(request))) + ) + ) + if request is None: + return tuple( + self.session.scalars( + self._filter_removed_task_producers( + select(TaskInstance), + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + region_id=region_id, + region_index=region_index, + map_indexes=map_indexes, + ).order_by(TaskInstance.region_index) + ) + ) + return resolve_current_producers( + dag_id=request.dag_id, + run_id=request.run_id, + task_id=request.task_id, + is_mapped=request.is_mapped, + context=request.context, + map_indexes=request.map_indexes, + region_id=request.region_id, + region_index=request.region_index, + session=self.session, + ) + + def select_skip_target_ids( + self, *, caller: TaskInstance, task_id: str, map_indexes: int | None = None + ) -> Select[tuple[UUID]]: + """ + Select the task instance ids a skip request from ``caller`` reaches for ``task_id``. + + A caller outside the loop that encloses ``task_id`` reaches every live pass and region of it. + """ + try: + task = self.get_task(caller.dag_id, caller.run_id, task_id, dag_version_id=caller.dag_version_id) + except TaskNotFound: + task = None + loop = enclosing_loop(task) if task is not None else None + if loop is not None: + caller_loop = enclosing_loop( + self.get_task( + caller.dag_id, caller.run_id, caller.task_id, dag_version_id=caller.dag_version_id + ) + ) + if caller_loop is None or caller_loop.group_id != loop.group_id: + query = select(TaskInstance.id).where( + TaskInstance.working_set.is_(True), + TaskInstance.dag_id == caller.dag_id, + TaskInstance.run_id == caller.run_id, + TaskInstance.task_id == task_id, + ) + if map_indexes is not None: + query = query.where(public_map_index_expression(TaskInstance) == map_indexes) + return query + return self.select_producer_ids( + dag_id=caller.dag_id, + run_id=caller.run_id, + task_id=task_id, + caller=caller, + map_indexes=map_indexes, + ) + + def _select_loop_producer_ids(self, request: _ProducerRequest) -> Select[tuple[UUID]]: + if TYPE_CHECKING: + assert request.loop_node_id is not None + return select_loop_producer_ids( + dag_id=request.dag_id, + run_id=request.run_id, + loop_node_id=request.loop_node_id, + task_id=request.task_id, + is_mapped=request.is_mapped, + session=self.session, + ) + + def select_producer_ids( + self, + *, + dag_id: str, + run_id: str, + task_id: str, + caller: TaskInstance | None = None, + region_id: UUID | None = None, + region_index: int | None = None, + map_indexes: int | Collection[int] | None = None, + previous_iteration: bool = False, + all_iterations: bool = False, + ) -> Select[tuple[UUID]]: + """Select the task instance ids :meth:`resolve` returns, without loading the producers.""" + if region_index is not None and region_id is None: + raise ValueError("region_index requires an explicit producer region_id") + if self._is_legacy_lookup(dag_id, run_id, task_id, region_id, previous_iteration): + return self._filter_legacy_producers( + select(TaskInstance.id), + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + region_index=region_index, + map_indexes=map_indexes, + ) + request = self._build_producer_request( + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + caller=caller, + region_id=region_id, + region_index=region_index, + map_indexes=map_indexes, + previous_iteration=previous_iteration, + all_iterations=all_iterations, + ) + if request is not None and request.loop_node_id is not None: + return self._select_loop_producer_ids(request) + if request is None: + return self._filter_removed_task_producers( + select(TaskInstance.id), + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + region_id=region_id, + region_index=region_index, + map_indexes=map_indexes, + ) + return select_current_producer_ids( + dag_id=request.dag_id, + run_id=request.run_id, + task_id=request.task_id, + is_mapped=request.is_mapped, + context=request.context, + map_indexes=request.map_indexes, + region_id=request.region_id, + region_index=request.region_index, + session=self.session, ) - if region_id is not None: - query = query.where(TaskInstance.region_id == region_id) - if region_index is not None: - query = query.where(TaskInstance.region_index == region_index) - if isinstance(map_indexes, int): - query = query.where(TaskInstance.region_index == map_indexes) - elif map_indexes is not None: - query = query.where(TaskInstance.region_index.in_(map_indexes)) - return tuple(self.session.scalars(query.order_by(TaskInstance.region_index))) diff --git a/airflow-core/src/airflow/serialization/definitions/xcom_arg.py b/airflow-core/src/airflow/serialization/definitions/xcom_arg.py index 4891448ebafdf..93327772a6545 100644 --- a/airflow-core/src/airflow/serialization/definitions/xcom_arg.py +++ b/airflow-core/src/airflow/serialization/definitions/xcom_arg.py @@ -22,13 +22,14 @@ from typing import TYPE_CHECKING, Any import attrs -from sqlalchemy import func, select +from sqlalchemy import func, or_ from sqlalchemy.orm import Session -from airflow.models.dynamic_region import ProducerContext, resolve_current_producers +from airflow.models.dynamic_region import ProducerContext, select_current_producer_ids from airflow.models.referencemixin import ReferenceMixin from airflow.models.xcom import XCOM_RETURN_KEY from airflow.serialization.definitions.notset import NOTSET, is_arg_set +from airflow.utils.db import exists_query from airflow.utils.state import State __all__ = ["SchedulerXComArg", "deserialize_xcom_arg", "get_task_map_length"] @@ -175,7 +176,7 @@ def _( if not mapped and xcom_arg.operator.get_needs_expansion(): return None - producers = resolve_current_producers( + producer_ids = select_current_producer_ids( dag_id=dag_id, run_id=run_id, task_id=task_id, @@ -183,14 +184,20 @@ def _( context=(producer_contexts or {}).get(task_id), session=session, ) - if not producers: - return None - if mapped and any(ti.state in State.unfinished for ti in producers): - return None + if mapped and exists_query( + TaskInstance.id.in_(producer_ids), + # 'state' can be NULL, which "NOT IN" would silently drop. + or_( + TaskInstance.state.is_(None), + TaskInstance.state.in_(s.value for s in State.unfinished if s is not None), + ), + session=session, + ): + return None # Not all of the expanded tis are done yet. read = XComModel.get_many( run_id=run_id, key=XCOM_RETURN_KEY, - producer_ids=select(TaskInstance.id).where(TaskInstance.id.in_([ti.id for ti in producers])), + producer_ids=producer_ids, ) entity = xcom_entity(read) diff --git a/airflow-core/src/airflow/ti_deps/dep_context.py b/airflow-core/src/airflow/ti_deps/dep_context.py index 43a9aee46078d..4083793c06730 100644 --- a/airflow-core/src/airflow/ti_deps/dep_context.py +++ b/airflow-core/src/airflow/ti_deps/dep_context.py @@ -169,12 +169,7 @@ def upstream_tis(self, ti: TaskInstance, task_id: str, *, session: Session) -> t task_id, ) if key not in self.producer_tis: - self.producer_tis[key] = resolver.resolve( - dag_id=ti.dag_id, - run_id=ti.run_id, - task_id=task_id, - caller=ti, - ) + self.producer_tis[key] = resolver.resolve_dependency(ti, task_id) return self.producer_tis[key] mapped_group_skip_decisions: dict[ diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 7c488791b93a6..2bb39e8692e90 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -18,8 +18,11 @@ from __future__ import annotations import re +import time +from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from datetime import datetime +from threading import Barrier, Event, local from types import SimpleNamespace from typing import TYPE_CHECKING from unittest import mock @@ -35,31 +38,36 @@ from opentelemetry.trace import StatusCode from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator from pydantic import ValidationError -from sqlalchemy import select, update +from sqlalchemy import delete, event, select, update from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session +from sqlalchemy.orm.exc import StaleDataError from airflow._shared.observability.traces import OverrideableRandomIdGenerator from airflow._shared.state import TaskScope from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator from airflow.api_fastapi.execution_api.app import lifespan +from airflow.api_fastapi.execution_api.datamodels.taskinstance import TISuccessStatePayload from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken -from airflow.api_fastapi.execution_api.routes.task_instances import _emit_task_span +from airflow.api_fastapi.execution_api.routes import task_instances as task_instances_route +from airflow.api_fastapi.execution_api.routes.task_instances import _emit_task_span, ti_update_state +from airflow.api_fastapi.execution_api.routes.xcoms import set_xcom from airflow.api_fastapi.execution_api.security import require_auth from airflow.exceptions import AirflowSkipException, TaskNotFound from airflow.jobs.scheduler_job_runner import SchedulerJobRunner -from airflow.models import RenderedTaskInstanceFields, TaskReschedule, Trigger +from airflow.models import DagRun, RenderedTaskInstanceFields, TaskReschedule, Trigger from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent, AssetModel from airflow.models.dag import DagModel +from airflow.models.dag_version import DagVersion from airflow.models.dagbag import DBDagBag from airflow.models.dynamic_region import DynamicRegion from airflow.models.log import Log -from airflow.models.task_coordinates import TaskCoordinateResolver +from airflow.models.task_coordinates import LOOP_GATE_OPERATOR, TaskCoordinateResolver from airflow.models.task_state_store import TaskStateStoreModel from airflow.models.taskinstance import TaskInstance, clear_task_instances -from airflow.models.xcom import XComModel +from airflow.models.xcom import XComModel, XComModelV2 from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.operators.python import PythonOperator from airflow.sdk import Asset, TaskGroup, TriggerRule, task, task_group @@ -399,6 +407,59 @@ def test_execution_breadcrumbs_keep_regional_identity_separate_from_map_index(cl class TestTIRunState: + @pytest.mark.parametrize("mapped", [False, True]) + @pytest.mark.parametrize("version", ["2026-06-30", "2026-10-30"]) + def test_startup_loop_context_uses_enclosing_iteration(self, client, session, dag_maker, mapped, version): + @task_group + def body(): + if mapped: + PythonOperator.partial(task_id="terminal", python_callable=str).expand(op_args=[[1], [2]]) + else: + EmptyOperator(task_id="terminal") + + with dag_maker(serialized=True): + loop = create_loop(body, max_iterations=4) + dr = dag_maker.create_dagrun(state=DagRunState.RUNNING) + ti = next(ti for ti in dr.task_instances if ti.task_id == loop.terminal_task_id) + parent = DynamicRegion.get_or_create( + dag_id=dr.dag_id, run_id=dr.run_id, node_id=loop.group_id, session=session + ) + session.add(parent) + session.flush() + ti.region_id, ti.region_index = parent.id, 2 + if mapped: + child = DynamicRegion.get_or_create( + dag_id=dr.dag_id, + run_id=dr.run_id, + node_id=ti.task_id, + parent_region_id=parent.id, + parent_region_index=2, + session=session, + ) + session.add(child) + session.flush() + ti.region_id, ti.region_index = child.id, 0 + ti.state = State.QUEUED + session.commit() + + response = client.patch( + f"/execution/task-instances/{ti.id}/run", + json=self.RUN_PAYLOAD, + headers={"Airflow-API-Version": version}, + ) + + assert response.status_code == 200 + if version == "2026-06-30": + assert "loop" not in response.json() + return + assert response.json()["loop"] == { + "node_id": loop.group_id, + "index": 2, + "max_iterations": 4, + "terminal_task_id": loop.terminal_task_id, + "terminal_is_mapped": mapped, + } + @pytest.mark.parametrize("regional", [False, True]) def test_startup_xcom_cleanup_keys_belong_to_the_starting_try( self, client, session, create_task_instance, regional @@ -1604,6 +1665,545 @@ def test_ti_run_creates_audit_log(self, client, session, create_task_instance, t class TestTIUpdateState: + def test_loop_continues_while_earlier_body_branch_is_running(self, client, session, dag_maker): + @task_group + def body(): + left = PythonOperator(task_id="left", python_callable=list) + right = PythonOperator(task_id="right", python_callable=list) + terminal = PythonOperator( + task_id="terminal", python_callable=list, trigger_rule=TriggerRule.ONE_SUCCESS + ) + [left, right] >> terminal + + with dag_maker(serialized=True, session=session): + loop = create_loop(body, max_iterations=2) + dr = dag_maker.create_dagrun() + dr.dag = dag_maker.serialized_dag + tis = {ti.task_id: ti for ti in dr.task_instances} + tis["body.left"].state = State.SUCCESS + right = tis["body.right"] + right.state = State.RUNNING + right_id = right.id + terminal = tis["body.terminal"] + gate = tis[loop.gate_task_id] + session.flush() + assert terminal in dr.task_instance_scheduling_decisions(session=session).schedulable_tis + terminal.state, terminal.start_date = State.RUNNING, DEFAULT_START_DATE + session.commit() + assert ( + client.patch( + f"/execution/task-instances/{terminal.id}/state", + json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()}, + ).status_code + == 204 + ) + session.expire_all() + assert gate in dr.task_instance_scheduling_decisions(session=session).schedulable_tis + gate.state, gate.start_date = State.RUNNING, DEFAULT_START_DATE + session.commit() + exec_app = client.app.routes[-1].app + exec_app.dependency_overrides[require_auth] = lambda: TIToken(id=gate.id, claims=TIClaims()) + assert ( + client.post( + f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{gate.task_id}/_airflow_loop_decision", + json="continue", + ).status_code + == 201 + ) + + response = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()}, + ) + + assert response.status_code == 204, response.text + session.expire_all() + assert session.get(TaskInstance, right_id).state == State.RUNNING + ready = dr.task_instance_scheduling_decisions(session=session).schedulable_tis + assert {(ti.task_id, ti.region_index) for ti in ready} == {("body.left", 1), ("body.right", 1)} + + @conf_vars({("state_store", "clear_on_success"): "true"}) + @pytest.mark.parametrize("mapped", [False, True]) + @pytest.mark.parametrize( + ("decision", "state", "max_iterations", "conditional", "rejected", "passes"), + [ + ("continue", State.SUCCESS, 3, False, False, [0, 1]), + ("stop", State.SUCCESS, 3, False, True, [0]), + (None, State.SUCCESS, 3, False, True, [0]), + ("continue", State.FAILED, 3, False, False, [0]), + ("continue", State.SKIPPED, 3, False, False, [0]), + ("stop", State.SUCCESS, 1, False, False, [0]), + ("continue", State.SUCCESS, 1, False, True, [0]), + ("stop", State.SUCCESS, 3, True, False, [0]), + ("stop", State.SUCCESS, 1, True, False, [0]), + ("continue", State.SUCCESS, 1, True, True, [0]), + (None, State.FAILED, 1, True, False, [0]), + ], + ) + def test_loop_gate_completion_consumes_decision_atomically( + self, + client, + session, + dag_maker, + decision, + state, + max_iterations, + conditional, + mapped, + rejected, + passes, + mocker, + ): + backend = mocker.create_autospec(MetastoreBackend, instance=True) + mocker.patch( + "airflow.api_fastapi.execution_api.routes.task_instances.get_state_backend", + autospec=True, + return_value=backend, + ) + + @task_group + def body(): + if mapped: + PythonOperator.partial(task_id="terminal", python_callable=list).expand(op_kwargs=[{}, {}]) + else: + PythonOperator(task_id="terminal", python_callable=list) + + with dag_maker(serialized=True, session=session): + loop = create_loop( + body, max_iterations=max_iterations, until=(lambda loop: True) if conditional else None + ) + dr = dag_maker.create_dagrun() + gate = next(ti for ti in dr.task_instances if ti.task_id == loop.gate_task_id) + gate.state = State.RUNNING + gate.start_date = DEFAULT_START_DATE + session.commit() + if decision and rejected: + XComModel.set_for_attempt( + task_instance_id=gate.id, + key="_airflow_loop_decision", + value=decision, + serialize=False, + session=session, + ) + session.commit() + elif decision: + exec_app = client.app.routes[-1].app + exec_app.dependency_overrides[require_auth] = lambda: TIToken(id=gate.id, claims=TIClaims()) + response = client.post( + f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{gate.task_id}/_airflow_loop_decision", + json=decision, + ) + assert response.status_code == 201 + + response = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()}, + ) + + assert response.status_code == 204, response.text + assert backend.clear.call_count == (state == State.SUCCESS and not rejected) + session.expire_all() + assert gate.state == (State.FAILED if rejected else state) + current = dr.get_task_instances(session=session) + assert sorted(ti.region_index for ti in current if ti.task_id == gate.task_id) == passes + if mapped: + children = session.scalars( + select(DynamicRegion).where(DynamicRegion.parent_region_id == gate.region_id) + ).all() + assert sorted(region.parent_region_index for region in children) == passes + region_ids = {gate.region_id, *(region.id for region in children)} + assert all(ti.region_id in region_ids for ti in current) + else: + assert all(ti.region_id == gate.region_id for ti in current) + signal = session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id)) + assert signal is None + if not rejected: + duplicate = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()}, + ) + assert duplicate.status_code == 200 + assert ( + sorted( + ti.region_index + for ti in dr.get_task_instances(session=session) + if ti.task_id == gate.task_id + ) + == passes + ) + if mapped and passes == [0, 1]: + successor_region = next(region for region in children if region.parent_region_index == 1) + assert [ti.region_index for ti in current if ti.region_id == successor_region.id] == [-1] + dr.dag = dag_maker.serialized_dag + + dr.task_instance_scheduling_decisions(session=session) + + assert sorted( + ti.region_index + for ti in dr.get_task_instances(session=session) + if ti.region_id == successor_region.id + ) == [0, 1] + + @pytest.fixture + def running_loop_gate(self, session, dag_maker): + @task_group + def body(): + PythonOperator(task_id="terminal", python_callable=list) + + with dag_maker(serialized=True, session=session): + loop = create_loop(body, max_iterations=3) + dr = dag_maker.create_dagrun() + gate = next(ti for ti in dr.task_instances if ti.task_id == loop.gate_task_id) + gate.state, gate.start_date = State.RUNNING, DEFAULT_START_DATE + XComModel.set_for_attempt( + task_instance_id=gate.id, + key="_airflow_loop_decision", + value="continue", + serialize=False, + session=session, + ) + session.commit() + return gate + + def test_loop_gate_materialization_error_rolls_back_state_signal_and_successor( + self, client, session, running_loop_gate, mocker + ): + gate = running_loop_gate + dr = gate.dag_run + original = DagRun._create_task_instances + observed = [] + + def fail_after_insert(*args, **kwargs): + original(*args, **kwargs) + with Session(bind=session.get_bind()) as observer: + observed.append(observer.get(TaskInstance, gate.id).state) + observed.append( + observer.scalars( + select(TaskInstance.region_index).where( + TaskInstance.dag_id == dr.dag_id, TaskInstance.task_id == gate.task_id + ) + ).all() + ) + raise StaleDataError("injected materialization failure") + + mocker.patch.object(DagRun, "_create_task_instances", autospec=True, side_effect=fail_after_insert) + response = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()}, + ) + + assert response.status_code == 500 + session.expire_all() + assert observed == [State.RUNNING, [0]] + assert gate.state == State.RUNNING + assert len(dr.get_task_instances(session=session)) == 2 + assert session.scalar(select(XComModel.value).where(XComModel.task_id == gate.task_id)) == "continue" + + @pytest.mark.parametrize( + ("max_tries", "expected_state"), + [(0, State.FAILED), (1, State.UP_FOR_RETRY)], + ) + def test_loop_gate_with_invalid_decision_leaves_running_in_same_request( + self, client, session, running_loop_gate, max_tries, expected_state + ): + gate = running_loop_gate + dr = gate.dag_run + gate_id = gate.id + gate.max_tries = max_tries + session.execute(delete(XComModelV2).where(XComModelV2.task_instance_id == gate.id)) + session.commit() + + response = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()}, + ) + + assert response.status_code == 204 + session.expire_all() + current = dr.get_task_instances(session=session) + assert sorted(ti.state for ti in current if ti.task_id == gate.task_id) == [expected_state] + assert len(current) == 2 + assert "requires a decision" in session.get(TaskInstance, gate_id).retry_reason + + @pytest.mark.backend("mysql", "postgres") + def test_concurrent_loop_gate_completions_create_one_successor(self, session, running_loop_gate): + gate = running_loop_gate + dr = gate.dag_run + gate_id = gate.id + bind = session.get_bind() + barrier = Barrier(2) + + def complete(): + with Session(bind=bind) as request_session: + barrier.wait(timeout=10) + result = ti_update_state( + task_instance_id=gate_id, + ti_patch_payload=TISuccessStatePayload(state="success", end_date=DEFAULT_END_DATE), + session=request_session, + dag_bag=DBDagBag(), + ) + return result.status_code if result is not None else 204 + + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(complete) for _ in range(2)] + assert sorted(future.result(timeout=20) for future in futures) == [200, 204] + session.expire_all() + assert gate.state == State.SUCCESS + assert sorted( + ti.region_index for ti in dr.get_task_instances(session=session) if ti.task_id == gate.task_id + ) == [0, 1] + assert session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id)) is None + + @pytest.mark.backend("mysql", "postgres") + def test_loop_decision_rewrite_and_completion_use_same_lock_order( + self, session, running_loop_gate, request + ): + gate = running_loop_gate + gate_id, dag_id, run_id, task_id = gate.id, gate.dag_id, gate.run_id, gate.task_id + XComModel.set_for_attempt( + task_instance_id=gate_id, + key="_airflow_loop_decision", + value="stop", + serialize=False, + session=session, + ) + session.commit() + bind = session.get_bind() + role = local() + writer_locked = Event() + completion_at_lock = Event() + + def note_completion_lock(conn, cursor, statement, parameters, context, executemany): + if getattr(role, "name", None) == "completion" and "FOR UPDATE" in statement: + completion_at_lock.set() + + def pause_writer_after_first_lock(conn, cursor, statement, parameters, context, executemany): + if ( + getattr(role, "name", None) == "writer" + and "FOR UPDATE" in statement + and not writer_locked.is_set() + ): + writer_locked.set() + assert completion_at_lock.wait(timeout=10) + # Give completion time to take its first lock, so a writer that locks in a different + # order deadlocks against it. + time.sleep(0.5) + + event.listen(bind, "before_cursor_execute", note_completion_lock) + event.listen(bind, "after_cursor_execute", pause_writer_after_first_lock) + request.addfinalizer(lambda: event.remove(bind, "before_cursor_execute", note_completion_lock)) + request.addfinalizer( + lambda: event.remove(bind, "after_cursor_execute", pause_writer_after_first_lock) + ) + + def rewrite(): + role.name = "writer" + with Session(bind=bind) as request_session: + set_xcom( + dag_id=dag_id, + run_id=run_id, + task_id=task_id, + key="_airflow_loop_decision", + session=request_session, + dag_bag=DBDagBag(), + value="continue", + token=TIToken(id=gate_id, claims=TIClaims()), + ) + request_session.commit() + + def complete(): + role.name = "completion" + assert writer_locked.wait(timeout=10) + with Session(bind=bind) as request_session: + ti_update_state( + task_instance_id=gate_id, + ti_patch_payload=TISuccessStatePayload(state="success", end_date=DEFAULT_END_DATE), + session=request_session, + dag_bag=DBDagBag(), + ) + + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(rewrite), pool.submit(complete)] + for future in futures: + future.result(timeout=20) + session.expire_all() + assert gate.state == State.SUCCESS + assert session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id)) is None + assert sorted( + ti.region_index + for ti in gate.dag_run.get_task_instances(session=session) + if ti.task_id == gate.task_id + ) == [0, 1] + + def test_loop_gate_completion_uses_its_pinned_limit(self, client, session, running_loop_gate, dag_maker): + gate = running_loop_gate + original_version = gate.dag_version_id + + @task_group + def body(): + PythonOperator(task_id="terminal", python_callable=list) + + with dag_maker(serialized=True, session=session): + create_loop(body, max_iterations=1) + + response = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()}, + ) + + assert response.status_code == 204, response.text + session.expire_all() + successors = [ti for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 1] + assert len(successors) == 2 + assert all(ti.dag_version_id == original_version for ti in successors) + + @pytest.mark.parametrize("state", ["success", "failed"]) + def test_loop_gate_missing_from_its_pinned_version_fails( + self, client, session, running_loop_gate, dag_maker, state + ): + gate = running_loop_gate + + def renamed_gate(loop): + return True + + @task_group + def body(): + PythonOperator(task_id="terminal", python_callable=list) + + with dag_maker(serialized=True, session=session): + create_loop(body, max_iterations=3, until=renamed_gate) + gate.dag_version_id = DagVersion.get_latest_version(gate.dag_id, session=session).id + session.commit() + + response = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()}, + ) + + assert response.status_code == 204, response.text + session.expire_all() + assert gate.state == State.FAILED + assert [ti.region_index for ti in gate.dag_run.get_task_instances(session=session)] == [0, 0] + + def test_late_loop_decision_cannot_recreate_consumed_signal(self, client, session, running_loop_gate): + gate = running_loop_gate + response = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()}, + ) + assert response.status_code == 204 + exec_app = client.app.routes[-1].app + exec_app.dependency_overrides[require_auth] = lambda: TIToken(id=gate.id, claims=TIClaims()) + + response = client.post( + f"/execution/xcoms/{gate.dag_id}/{gate.run_id}/{gate.task_id}/_airflow_loop_decision", + json="continue", + ) + + assert response.status_code == 409 + assert session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id)) is None + + def test_loop_gate_state_update_error_rolls_back_next_pass_and_keeps_decision( + self, client, session, running_loop_gate, mocker + ): + gate = running_loop_gate + mocker.patch.object( + task_instances_route, + "_create_ti_state_update_query_and_update_state", + autospec=True, + side_effect=StaleDataError("injected state update failure"), + ) + + response = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()}, + ) + + assert response.status_code == 500 + session.expire_all() + assert gate.state == State.RUNNING + assert len(gate.dag_run.get_task_instances(session=session)) == 2 + assert session.scalar(select(XComModel.value).where(XComModel.task_id == gate.task_id)) == "continue" + + @pytest.mark.parametrize("state", ["success", "failed"]) + def test_non_loop_operator_named_loop_gate_completes_normally( + self, client, session, create_task_instance, state + ): + ti = create_task_instance(task_id="user_gate", start_date=DEFAULT_START_DATE, state=State.RUNNING) + ti.operator = LOOP_GATE_OPERATOR + session.commit() + + response = client.patch( + f"/execution/task-instances/{ti.id}/state", + json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()}, + ) + + assert response.status_code == 204 + session.expire_all() + assert ti.state == state + + def test_loop_gate_inner_insert_error_is_not_swallowed(self, client, session, running_loop_gate, mocker): + gate = running_loop_gate + bulk_insert = mocker.patch.object( + Session, + "bulk_insert_mappings", + autospec=True, + side_effect=StaleDataError("injected insert error"), + ) + + response = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()}, + ) + + assert response.status_code == 500 + bulk_insert.assert_called_once() + session.expire_all() + assert gate.state == State.RUNNING + assert len(gate.dag_run.get_task_instances(session=session)) == 2 + assert session.scalar(select(XComModel.value).where(XComModel.task_id == gate.task_id)) == "continue" + + @pytest.mark.parametrize("state", [State.SUCCESS, State.FAILED, State.SKIPPED]) + def test_rerun_gate_preserves_existing_later_pass(self, client, session, running_loop_gate, state): + gate = running_loop_gate + response = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()}, + ) + assert response.status_code == 204 + session.refresh(gate) + successor_ids = { + ti.id for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 1 + } + archived_id = gate.id + dag_run = gate.dag_run + clear_task_instances([gate], session=session) + session.expire_all() + gate = session.scalar( + select(TaskInstance).where( + TaskInstance.dag_id == dag_run.dag_id, + TaskInstance.run_id == dag_run.run_id, + TaskInstance.task_id == gate.task_id, + TaskInstance.region_index == 0, + TaskInstance.working_set.is_(True), + ) + ) + assert gate.id != archived_id + gate.state = State.RUNNING + session.commit() + + response = client.patch( + f"/execution/task-instances/{gate.id}/state", + json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()}, + ) + + assert response.status_code == 204, response.text + session.expire_all() + assert gate.state == state + assert { + ti.id for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 1 + } == successor_ids + @pytest.mark.parametrize( ("interruption", "expected_state", "expected_status"), [("clear", State.RESTARTING, 409), ("failed", State.FAILED, 409)], @@ -3459,6 +4059,38 @@ def body(): assert target.state == State.SKIPPED assert previous.state is None + @pytest.mark.parametrize("tasks", [["body.target"], [["body.target", -1]]]) + def test_skip_downstream_from_outside_loop_skips_every_live_pass(self, client, session, dag_maker, tasks): + @task_group + def body(): + EmptyOperator(task_id="target") + + with dag_maker(serialized=True) as dag: + EmptyOperator(task_id="branch") >> create_loop(body, max_iterations=3) + dr = dag_maker.create_dagrun() + region = DynamicRegion.get_or_create( + dag_id=dr.dag_id, run_id=dr.run_id, node_id="body", session=session + ) + session.add(region) + session.flush() + tis = {ti.task_id: ti for ti in dr.task_instances} + caller, target = tis["branch"], tis["body.target"] + other_pass = TaskInstance( + task=dag.get_task(target.task_id), run_id=dr.run_id, dag_version_id=target.dag_version_id + ) + other_pass.region_id, other_pass.region_index = region.id, 1 + session.add(other_pass) + session.commit() + + response = client.patch( + f"/execution/task-instances/{caller.id}/skip-downstream", json={"tasks": tasks} + ) + + assert response.status_code == 204 + session.expire_all() + assert target.state == State.SKIPPED + assert other_pass.state == State.SKIPPED + def setup_method(self): clear_db_runs() diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py index a3e95dca3a566..5c8d49c3f1e5b 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py @@ -25,7 +25,8 @@ import pytest from fastapi import FastAPI, HTTPException, Path, Request, status -from sqlalchemy import delete, select, update +from sqlalchemy import delete, event, select, update +from sqlalchemy.orm import Session from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import JWTValidator @@ -129,6 +130,478 @@ def body(): return dr, producer, consumer, previous +@pytest.mark.parametrize("suffix", ["", "/item/0", "/item/-1", "/slice"]) +def test_previous_iteration_xcom_uses_retained_predecessor(client, loop_xcoms, suffix): + dr, producer, _, _ = loop_xcoms + url = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{producer.task_id}/key" + params = {"previous_iteration": True} + + response = client.get(url + suffix, params=params) + + assert response.status_code == 200 + expected = {"": {"key": "key", "value": "previous"}, "/slice": ["previous"]} + assert response.json() == expected.get(suffix, "previous") + assert client.head(url, params=params).headers["Content-Range"] == "map_indexes 1" + + +def test_previous_iteration_rejects_prior_dates(client, loop_xcoms): + dr, producer, _, _ = loop_xcoms + + response = client.get( + f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{producer.task_id}/key", + params={"previous_iteration": True, "include_prior_dates": True}, + ) + + assert response.status_code == 400 + + +def test_previous_mapped_iteration_reads_retained_slots_before_count_and_slice( + client, dag_maker, session, authenticate_as +): + @task_group + def body(): + PythonOperator.partial(task_id="mapped", python_callable=str).expand( + op_args=[[1], [2]] + ) >> EmptyOperator(task_id="consumer") + + with dag_maker(serialized=True) as dag: + create_loop(body, max_iterations=3) + dr = dag_maker.create_dagrun() + parent = DynamicRegion.get_or_create(dag_id=dr.dag_id, run_id=dr.run_id, node_id="body", session=session) + session.add(parent) + session.flush() + previous = DynamicRegion.get_or_create( + dag_id=dr.dag_id, + run_id=dr.run_id, + node_id="body.mapped", + parent_region_id=parent.id, + parent_region_index=0, + session=session, + ) + current = DynamicRegion.get_or_create( + dag_id=dr.dag_id, + run_id=dr.run_id, + node_id="body.mapped", + parent_region_id=parent.id, + parent_region_index=1, + session=session, + ) + session.add_all([previous, current]) + session.flush() + replacement = DynamicRegion( + dag_id=dr.dag_id, + run_id=dr.run_id, + node_id="body.mapped", + parent_region_id=parent.id, + parent_region_index=0, + forked_from_region_id=previous.id, + ) + session.add(replacement) + session.flush() + consumer = next(ti for ti in dr.task_instances if ti.task_id == "body.consumer") + consumer.region_id, consumer.region_index = parent.id, 1 + for ti in dr.task_instances: + if ti.task_id == "body.mapped": + ti.region_id = current.id + producers = {} + for region, index in [(previous, 0), (previous, 1), (replacement, 1)]: + ti = TaskInstance( + task=dag.get_task("body.mapped"), run_id=dr.run_id, dag_version_id=consumer.dag_version_id + ) + ti.region_id, ti.region_index = region.id, index + session.add(ti) + producers[region.id, index] = ti + session.flush() + producers[previous.id, 1].archive(reason="cleared", session=session) + for (region, index), value in [ + ((previous, 0), "retained"), + ((previous, 1), "archived"), + ((replacement, 1), "replacement"), + ]: + XComModel.set_for_attempt( + task_instance_id=producers[region.id, index].id, + key="key", + value=value, + serialize=False, + session=session, + ) + session.commit() + authenticate_as(consumer) + url = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/body.mapped/key" + params = {"previous_iteration": True} + assert client.head(url, params=params).headers["Content-Range"] == "map_indexes 2" + assert client.get(url + "/slice", params=params).json() == ["retained", "replacement"] + assert client.get(url + "/item/-1", params=params).json() == "replacement" + + +LOOP_SEQUENCES = { + "unmapped": ("body.step", ["step-0", "step-1", "step-2"]), + "mapped": ("body.fanout", ["0-0", "0-1", "1-0", "1-1", "2-0", "2-1"]), +} + + +@pytest.fixture +def outside_loop_run(dag_maker, session, authenticate_as): + @task_group + def body(): + EmptyOperator(task_id="step") >> PythonOperator.partial(task_id="fanout", python_callable=str).expand( + op_args=[[1], [2]] + ) + + with dag_maker(serialized=True) as dag: + loop = create_loop(body, max_iterations=3) + PythonOperator( + task_id="outside", + python_callable=list, + op_args=[loop["step"].output, loop["fanout"].output], + ) + dr = dag_maker.create_dagrun() + root = session.scalar( + select(DynamicRegion).where( + DynamicRegion.dag_id == dr.dag_id, + DynamicRegion.run_id == dr.run_id, + DynamicRegion.node_id == "body", + ) + ) + tis = {(ti.task_id, ti.region_id, ti.region_index): ti for ti in dr.task_instances} + steps = {0: tis["body.step", root.id, 0]} + fanouts = {(0, ti.region_index): ti for ti in dr.task_instances if ti.task_id == "body.fanout"} + for iteration in (1, 2): + steps[iteration] = TaskInstance( + dag.get_task("body.step"), + dr.created_dag_version_id, + run_id=dr.run_id, + region_id=root.id, + region_index=iteration, + ) + nested = DynamicRegion.get_or_create( + dag_id=dr.dag_id, + run_id=dr.run_id, + node_id="body.fanout", + parent_region_id=root.id, + parent_region_index=iteration, + session=session, + ) + session.add_all([steps[iteration], nested]) + session.flush() + for slot in range(2): + fanouts[iteration, slot] = TaskInstance( + dag.get_task("body.fanout"), + dr.created_dag_version_id, + run_id=dr.run_id, + region_id=nested.id, + region_index=slot, + ) + session.add(fanouts[iteration, slot]) + session.flush() + for iteration in (2, 1, 0): + XComModel.set_for_attempt( + task_instance_id=steps[iteration].id, + key="key", + value=f"step-{iteration}", + serialize=False, + session=session, + ) + XComModel.set_for_attempt( + task_instance_id=steps[iteration].id, + key=XCOM_RETURN_KEY, + value=f"returned-{iteration}", + serialize=False, + session=session, + ) + for slot in (1, 0): + XComModel.set_for_attempt( + task_instance_id=fanouts[iteration, slot].id, + key="key", + value=f"{iteration}-{slot}", + serialize=False, + session=session, + ) + session.commit() + outside = next(ti for ti in dr.task_instances if ti.task_id == "outside") + authenticate_as(outside) + return dr, dag, root, steps, authenticate_as + + +def get_loop_url(dr, task_id): + return f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{task_id}/key" + + +@pytest.mark.parametrize("shape", LOOP_SEQUENCES) +def test_outside_caller_reads_every_iteration_of_a_loop_task(client, outside_loop_run, shape): + dr = outside_loop_run[0] + task_id, expected = LOOP_SEQUENCES[shape] + url = get_loop_url(dr, task_id) + + assert client.head(url).headers["Content-Range"] == f"map_indexes {len(expected)}" + assert client.get(f"{url}/slice").json() == expected + for offset, value in enumerate(expected): + assert client.get(f"{url}/item/{offset}").json() == value + assert client.get(f"{url}/item/{offset - len(expected)}").json() == value + assert client.get(f"{url}/item/{len(expected)}").status_code == 404 + assert client.get(url, params={"offset": -1}).json() == {"key": "key", "value": expected[-1]} + + +def test_outside_task_that_uses_loop_outputs_waits_for_the_gate_and_reads_the_loop_task_xcom( + client, dag_maker, outside_loop_run +): + dr = outside_loop_run[0] + consumer = dag_maker.serialized_dag.get_task("outside") + + assert consumer.upstream_task_ids == {"body.__loop_gate"} + url = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/body.step/{XCOM_RETURN_KEY}" + assert client.head(url).headers["Content-Range"] == "map_indexes 3" + assert client.get(f"{url}/slice").json() == ["returned-0", "returned-1", "returned-2"] + assert client.get(f"{url}/item/-1").json() == "returned-2" + + +@pytest.mark.parametrize("shape", LOOP_SEQUENCES) +@pytest.mark.parametrize( + "params", + [ + {"start": 1}, + {"stop": 3}, + {"stop": -1}, + {"start": -2}, + {"start": 1, "stop": -1}, + {"start": -3, "stop": 2}, + {"step": -1}, + {"step": 2}, + ], +) +def test_outside_caller_slices_a_loop_task_depth_first(client, outside_loop_run, shape, params): + dr = outside_loop_run[0] + task_id, expected = LOOP_SEQUENCES[shape] + + response = client.get(f"{get_loop_url(dr, task_id)}/slice", params=params) + + assert response.json() == expected[slice(params.get("start"), params.get("stop"), params.get("step"))] + + +@pytest.mark.parametrize("shape", LOOP_SEQUENCES) +def test_outside_caller_reads_a_forked_loop_from_the_region_that_owns_each_pass( + client, session, outside_loop_run, shape +): + dr, dag, root, steps, _ = outside_loop_run + fork = DynamicRegion( + dag_id=dr.dag_id, + run_id=dr.run_id, + node_id="body", + forked_from_region_id=root.id, + resumes_from_index=2, + ) + nested = DynamicRegion( + dag_id=dr.dag_id, + run_id=dr.run_id, + node_id="body.fanout", + parent_region_id=None, + parent_region_index=None, + ) + session.add(fork) + session.flush() + nested.parent_region_id, nested.parent_region_index = fork.id, 2 + session.add(nested) + session.flush() + replacements = [ + TaskInstance( + dag.get_task("body.step"), + dr.created_dag_version_id, + run_id=dr.run_id, + region_id=fork.id, + region_index=2, + ), + *( + TaskInstance( + dag.get_task("body.fanout"), + dr.created_dag_version_id, + run_id=dr.run_id, + region_id=nested.id, + region_index=slot, + ) + for slot in range(2) + ), + ] + session.add_all(replacements) + session.flush() + for ti, value in zip(replacements, ["step-2-again", "2-0-again", "2-1-again"]): + XComModel.set_for_attempt( + task_instance_id=ti.id, key="key", value=value, serialize=False, session=session + ) + session.commit() + task_id = LOOP_SEQUENCES[shape][0] + expected = { + "unmapped": ["step-0", "step-1", "step-2-again"], + "mapped": ["0-0", "0-1", "1-0", "1-1", "2-0-again", "2-1-again"], + }[shape] + url = get_loop_url(dr, task_id) + + assert client.head(url).headers["Content-Range"] == f"map_indexes {len(expected)}" + assert client.get(f"{url}/slice").json() == expected + assert client.get(f"{url}/item/-1").json() == expected[-1] + + +def test_outside_caller_cannot_read_a_single_value_of_a_loop_task(client, outside_loop_run): + dr = outside_loop_run[0] + + assert client.get(get_loop_url(dr, "body.step")).status_code == 400 + + +def test_outside_caller_cannot_read_the_previous_iteration_of_a_loop_task(client, outside_loop_run): + dr = outside_loop_run[0] + + response = client.get(f"{get_loop_url(dr, 'body.step')}/slice", params={"previous_iteration": True}) + + assert response.status_code == 400 + + +def test_outside_caller_cannot_read_a_loop_with_several_executions(client, session, outside_loop_run): + dr, _, root, _, _ = outside_loop_run + session.add( + DynamicRegion( + dag_id=dr.dag_id, + run_id=dr.run_id, + node_id="body", + parent_region_id=root.id, + parent_region_index=0, + ) + ) + session.commit() + + response = client.get(f"{get_loop_url(dr, 'body.step')}/slice") + + assert response.status_code == 400 + assert "several executions" in response.json()["detail"] + + +def test_loop_task_reads_from_inside_the_loop_stay_on_one_pass(client, outside_loop_run): + dr, _, _, steps, authenticate_as = outside_loop_run + authenticate_as(steps[1]) + step_url = get_loop_url(dr, "body.step") + fanout_url = get_loop_url(dr, "body.fanout") + + assert client.head(step_url).headers["Content-Range"] == "map_indexes 1" + assert client.get(f"{step_url}/slice").json() == ["step-1"] + assert client.get(f"{step_url}/slice", params={"previous_iteration": True}).json() == ["step-0"] + assert client.get(f"{fanout_url}/slice").json() == ["1-0", "1-1"] + assert client.get(f"{fanout_url}/item/-1").json() == "1-1" + + +@pytest.mark.parametrize("key", ["_airflow_loop_decision", "_airflow_loop_other"]) +def test_normal_xcom_push_rejects_loop_reserved_prefix(client, create_task_instance, key): + ti = create_task_instance() + response = client.post(f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/{key}", json="continue") + assert response.status_code == 400 + + +@pytest.mark.parametrize( + ("value", "region_index", "expected_status"), + [ + ("continue", 0, 201), + ("stop", 2, 201), + ("continue", 2, 400), + ("stop", 0, 400), + (True, 2, 400), + ("invalid", 2, 400), + ({"decision": "stop"}, 2, 400), + ], +) +def test_internal_loop_decision_uses_token_gate_coordinates( + client, loop_xcoms, session, authenticate_as, value, region_index, expected_status +): + dr, _, consumer, _ = loop_xcoms + gate = session.scalar( + select(TaskInstance).where( + TaskInstance.dag_id == dr.dag_id, TaskInstance.task_id == "body.__loop_gate" + ) + ) + gate.region_id, gate.region_index = consumer.region_id, region_index + gate.state = "running" + session.commit() + authenticate_as(gate) + url = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{gate.task_id}/_airflow_loop_decision" + + response = client.post(url, json=value) + + assert response.status_code == expected_status + stored = session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id)) + if expected_status == 400: + assert stored is None + return + assert (stored.region_id, stored.region_index, stored.value) == (gate.region_id, region_index, value) + authenticate_as(consumer) + assert client.post(url, json="stop").status_code == 400 + assert ( + client.post( + f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{consumer.task_id}/_airflow_loop_decision", + json="stop", + ).status_code + == 400 + ) + + +def test_loop_result_item_read_loads_no_producer_rows(dag_maker, session, client, authenticate_as): + @task_group + def body(): + PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1], [2], [3]]) >> ( + EmptyOperator(task_id="consumer") + ) + + with dag_maker(serialized=True) as dag: + create_loop(body, max_iterations=3) + dr = dag_maker.create_dagrun() + tis = {(ti.task_id, ti.map_index): ti for ti in dr.task_instances} + consumer = tis["body.consumer", 0] + loop_region = DynamicRegion.get_or_create( + dag_id=dr.dag_id, run_id=dr.run_id, node_id="body", session=session + ) + session.add(loop_region) + session.flush() + consumer.region_id, consumer.region_index = loop_region.id, 1 + mapped_task = dag.get_task("body.mapped") + for iteration in range(3): + region = DynamicRegion.get_or_create( + dag_id=dr.dag_id, + run_id=dr.run_id, + node_id="body.mapped", + parent_region_id=loop_region.id, + parent_region_index=iteration, + session=session, + ) + session.add(region) + session.flush() + for map_index in range(3): + ti = tis.get(("body.mapped", map_index)) if iteration == 0 else None + if ti is None: + ti = TaskInstance(mapped_task, consumer.dag_version_id, run_id=dr.run_id, map_index=map_index) + session.add(ti) + ti.region_id = region.id + session.flush() + XComModel.set_for_attempt( + task_instance_id=ti.id, + key="return_value", + value=f"{iteration}-{map_index}", + serialize=False, + session=session, + ) + session.commit() + authenticate_as(consumer) + session.expunge_all() + loaded = [] + + def record(_, instance): + if isinstance(instance, TaskInstance) and instance.task_id == "body.mapped": + loaded.append(instance.id) + + event.listen(Session, "loaded_as_persistent", record) + try: + response = client.get(f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/body.mapped/return_value/item/0") + finally: + event.remove(Session, "loaded_as_persistent", record) + + assert response.json() == "1-0" + assert loaded == [] + + def test_loop_xcom_omission_uses_consumer_pass_and_retained_region(client, loop_xcoms): dr, producer, consumer, _ = loop_xcoms base = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}" diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py index 2f4182505cb42..55280830e1496 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py @@ -49,6 +49,8 @@ def test_task_coordinate_schema_by_version(client, version): previous = response.json()["components"]["schemas"]["PreviousTIResponse"]["properties"] assert ("region_id" in previous) == (version == "2026-10-30") assert ("region_index" in previous) == (version == "2026-10-30") + context = response.json()["components"]["schemas"]["TIRunContext"]["properties"] + assert ("loop" in context) == (version == "2026-10-30") @pytest.mark.parametrize( diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index a66e3f6c5ca27..3203bd23dcda3 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -147,7 +147,9 @@ StartOfDayMapper, StartOfHourMapper, task, + task_group, ) +from airflow.sdk.definitions._internal.loop import create_loop from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback from airflow.sdk.definitions.timetables.assets import PartitionedAssetTimetable from airflow.serialization.definitions.dag import SerializedDAG @@ -1180,6 +1182,54 @@ def test_process_executor_events_raises_on_unknown_key_type(self, session, key): with pytest.raises(TypeError, match="Unknown workload key type in event buffer"): self.job_runner._process_executor_events(executor=executor, session=session) + @pytest.mark.parametrize("retries", [0, 2]) + def test_missing_gate_definition_executor_success_uses_failure_policy( + self, dag_maker, session, mocker, retries + ): + @task_group + def body(): + EmptyOperator(task_id="terminal") + + with dag_maker(serialized=True, default_args={"retries": retries}, session=session): + loop = create_loop(body, max_iterations=3) + dr = dag_maker.create_dagrun() + gate = next(ti for ti in dr.task_instances if ti.task_id == loop.gate_task_id) + executor = MockExecutor(do_update=False) + scheduler_job = Job() + session.add(scheduler_job) + session.flush() + runner = SchedulerJobRunner(scheduler_job, executors=[executor]) + mocker.patch.object(runner.scheduler_dag_bag, "get_dag_for_run", autospec=True, return_value=None) + gate.state = State.RUNNING + gate.try_number = 1 + gate.queued_by_job_id = scheduler_job.id + session.commit() + original_id = gate.id + callback = mocker.patch.object(executor, "send_callback", autospec=True) + executor.event_buffer[TaskInstanceUuid(gate.id)] = State.SUCCESS, None + + runner._process_executor_events(executor=executor, session=session) + + gate = session.scalar( + select(TaskInstance).where( + TaskInstance.dag_id == dr.dag_id, + TaskInstance.run_id == dr.run_id, + TaskInstance.task_id == loop.gate_task_id, + TaskInstance.working_set.is_(True), + ) + ) + assert gate.state == (State.UP_FOR_RETRY if retries else State.FAILED) + assert (gate.id != original_id) is bool(retries) + callback.assert_not_called() + if retries: + archived = session.scalar( + select(TaskInstance) + .where(TaskInstance.id == original_id) + .execution_options(include_all_attempts=True) + ) + assert archived.working_set is None + assert len(dr.get_task_instances(session=session)) == 2 + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") @mock.patch("airflow._shared.observability.metrics.stats._get_backend") def test_process_executor_event_missing_dag( diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 4b769b4d21439..47ad1c63955ad 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -55,6 +55,7 @@ from airflow._shared.timezones import timezone from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance as TIDataModel from airflow.callbacks.callback_requests import DagCallbackRequest, DagRunContext +from airflow.dag_processing.dagbag import DagBag from airflow.models.dag import DagModel, infer_automated_data_interval from airflow.models.dag_version import DagVersion from airflow.models.dagbag import DBDagBag @@ -85,12 +86,15 @@ task_group, teardown, ) +from airflow.sdk.definitions._internal.loop import create_loop from airflow.sdk.definitions.callback import AsyncCallback from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineReference, VariableInterval from airflow.serialization.definitions.deadline import SerializedReferenceModels from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG from airflow.settings import get_policy_plugin_manager from airflow.task.trigger_rule import TriggerRule +from airflow.ti_deps.dep_context import DepContext +from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep from airflow.triggers.base import StartTriggerArgs from airflow.utils.session import create_session from airflow.utils.sqlalchemy import prohibit_commit @@ -101,9 +105,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 +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 @@ -123,6 +128,258 @@ def get_callback_ti(ti): ) +@pytest.mark.parametrize("values", [[], [{}, {}]]) +def test_loop_first_pass_uses_one_region_and_mapped_child(dag_maker, session, values): + @task_group + def body(): + plain = PythonOperator(task_id="plain", python_callable=list) + mapped = PythonOperator.partial(task_id="mapped", python_callable=list).expand(op_kwargs=values) + plain >> mapped + + with dag_maker(serialized=True): + outside = EmptyOperator(task_id="outside") + loop = create_loop(body, max_iterations=3) + outside >> loop + dr = dag_maker.create_dagrun() + regions = session.scalars( + select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id, DynamicRegion.run_id == dr.run_id) + ).all() + assert {region.node_id for region in regions} == {loop.group_id, "body.mapped"} + parent = next(region for region in regions if region.node_id == loop.group_id) + mapped_region = next(region for region in regions if region.node_id == "body.mapped") + assert (mapped_region.parent_region_id, mapped_region.parent_region_index) == (parent.id, 0) + original_ids = {ti.id for ti in dr.task_instances} + for ti in dr.task_instances: + if ti.task_id == "outside": + assert (ti.region_id, ti.region_index) == (SENTINEL_REGION_ID, -1) + elif ti.task_id == "body.mapped": + assert ti.region_id == mapped_region.id + else: + assert (ti.region_id, ti.region_index) == (parent.id, 0) + + dr.verify_integrity(session=session, dag_version_id=dr.created_dag_version_id) + + assert {ti.id for ti in dr.get_task_instances(session=session)} == original_ids + assert ( + session.scalar( + select(func.count()).select_from(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id) + ) + == 2 + ) + + +@pytest.mark.parametrize( + ("trigger_rule", "left_state", "expected"), + [ + (TriggerRule.ALWAYS, State.RUNNING, True), + (TriggerRule.ONE_SUCCESS, State.SUCCESS, True), + (TriggerRule.ALL_SUCCESS, State.SUCCESS, False), + (TriggerRule.ALL_DONE, State.FAILED, False), + (TriggerRule.ALL_SUCCESS, State.SKIPPED, False), + ], +) +def test_loop_entry_uses_normal_trigger_rule_before_all_upstreams_finish( + dag_maker, session, trigger_rule, left_state, expected +): + @task_group + def body(): + PythonOperator(task_id="entry", python_callable=list, trigger_rule=trigger_rule) + + with dag_maker(serialized=True): + left = PythonOperator(task_id="left", python_callable=list) + right = PythonOperator(task_id="right", python_callable=list) + loop = create_loop(body, max_iterations=3) + [left, right] >> loop + dr = dag_maker.create_dagrun() + tis = {ti.task_id: ti for ti in dr.task_instances} + tis["left"].state = left_state + tis["right"].state = State.RUNNING + entry = tis["body.entry"] + entry.task = dag_maker.serialized_dag.get_task(entry.task_id) + assert entry.region_id != SENTINEL_REGION_ID + session.flush() + + assert TriggerRuleDep().is_met(ti=entry, dep_context=DepContext(), session=session) is expected + + +def test_loop_new_member_joins_live_pass_after_reserialization(dag_maker, session): + @task_group + def body(add_member=False): + terminal = PythonOperator(task_id="terminal", python_callable=list) + if add_member: + PythonOperator(task_id="new", python_callable=list) >> terminal + + with dag_maker(serialized=True, session=session): + loop = create_loop(body, max_iterations=3) + dr = dag_maker.create_dagrun() + existing = dr.get_task_instances(session=session) + original_ids = {ti.id for ti in existing} + gate = next(ti for ti in existing if ti.task_id == loop.gate_task_id) + gate.state = State.SUCCESS + session.flush() + + with dag_maker(serialized=True, session=session): + 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 + + dr.verify_integrity(session=session, dag_version_id=version_id) + + current = dr.get_task_instances(session=session) + assert {ti.id for ti in current if ti.task_id != "body.new"} == original_ids + added = next(ti for ti in current if ti.task_id == "body.new") + assert (added.region_id, added.region_index, added.dag_version_id) == (gate.region_id, 0, version_id) + assert session.get(TI, gate.id).state == State.SUCCESS + + +def test_loop_integrity_skips_loop_with_regions_but_no_live_gate(dag_maker, session, caplog): + @task_group + def body(): + PythonOperator(task_id="terminal", python_callable=list) + + with dag_maker(serialized=True, session=session): + EmptyOperator(task_id="outside") + loop = create_loop(body, max_iterations=3) + dr = dag_maker.create_dagrun() + for ti in dr.get_task_instances(session=session): + session.delete(ti) + session.flush() + + with caplog.at_level("WARNING"): + dr.verify_integrity(session=session, dag_version_id=dr.created_dag_version_id) + + assert [ti.task_id for ti in dr.get_task_instances(session=session)] == ["outside"] + assert { + "event": "skipping loop that has regions but no live gate", + "dag_id": dr.dag_id, + "run_id": dr.run_id, + "loop": loop.group_id, + "log_level": "warning", + } in caplog + assert ( + session.scalar( + select(func.count()).select_from(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id) + ) + == 1 + ) + + +def test_loop_integrity_reuses_region_created_by_a_concurrent_run(dag_maker, session): + @task_group + def body(): + PythonOperator(task_id="terminal", python_callable=list) + + with dag_maker(serialized=True, session=session): + loop = create_loop(body, max_iterations=3) + dr = dag_maker.create_dagrun() + region_id = session.scalar(select(DynamicRegion.id).where(DynamicRegion.dag_id == dr.dag_id)) + for ti in dr.get_task_instances(session=session): + session.delete(ti) + session.flush() + + with mock.patch.object(session, "scalar", autospec=True, return_value=None): + dr.verify_integrity(session=session, dag_version_id=dr.created_dag_version_id) + + gate = next(ti for ti in dr.get_task_instances(session=session) if ti.task_id == loop.gate_task_id) + assert gate.region_id == region_id + assert ( + session.scalar( + select(func.count()).select_from(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id) + ) + == 1 + ) + + +@pytest.mark.parametrize( + "mode", ["fixed", "mapped", "empty", "until_true", "until_false", "handled_failure", "unhandled_failure"] +) +def test_loop_runs_through_scheduler_sdk_and_api(dag_maker, session, mode): + @task + def terminal(value, *, ti, loop): + assert ti.region_id != SENTINEL_REGION_ID + if mode == "mapped": + assert ti.map_index == value - 1 + else: + assert ti.map_index == -1 + assert ti.region_index == loop.index + return value + + @task + def fail(): + raise ValueError("body branch failed") + + @task_group + def body(): + if mode in {"mapped", "empty"}: + terminal.expand(value=[1, 2] if mode == "mapped" else []) + elif mode == "handled_failure": + fail() >> terminal.override(trigger_rule=TriggerRule.ALL_DONE)(1) + elif mode == "unhandled_failure": + fail() >> terminal(1) + else: + terminal(1) + + def converged(*, loop): + return mode == "until_true" and loop.index == 1 + + with dag_maker(serialized=False) as dag: + loop = body.loop(max_iterations=2, until=converged if mode.startswith("until") else None) + loop >> PythonOperator(task_id="outside", python_callable=list, trigger_rule=TriggerRule.ALL_DONE) + + dr = dag.test() + + gates = sorted( + (ti.region_index, ti.state) + for ti in dr.get_task_instances(session=session) + if ti.task_id == loop.gate_task_id + ) + if mode == "empty": + assert gates == [(0, State.SKIPPED)] + elif mode == "until_false": + assert gates == [(0, State.SUCCESS), (1, State.FAILED)] + elif mode == "unhandled_failure": + assert gates == [(0, State.UPSTREAM_FAILED)] + else: + assert gates == [(0, State.SUCCESS), (1, State.SUCCESS)] + outside = dr.get_task_instance(task_id="outside", session=session) + assert outside.state == State.SUCCESS + + +@pytest.fixture +def clear_loop_example_runs(): + 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), + ("partial_override_loop", 3), + ], +) +def test_documented_task_loop_examples_complete(session, dag_id, expected_iterations): + bag = DagBag( + dag_folder=str(AIRFLOW_CORE_SOURCES_PATH / "airflow" / "example_dags" / "example_task_loops.py"), + ) + assert bag.import_errors == {} + dag = bag.get_dag(dag_id) + + dr = dag.test() + + assert dr.state == DagRunState.SUCCESS + 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) + ] + + async def empty_callback_for_deadline(): """Used in a number of tests to confirm that Deadlines and DeadlineAlerts function correctly.""" pass diff --git a/airflow-core/tests/unit/models/test_task_coordinates.py b/airflow-core/tests/unit/models/test_task_coordinates.py index 30131fe22770b..c872b653b4a16 100644 --- a/airflow-core/tests/unit/models/test_task_coordinates.py +++ b/airflow-core/tests/unit/models/test_task_coordinates.py @@ -19,6 +19,8 @@ import pytest from sqlalchemy import select +from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken +from airflow.api_fastapi.execution_api.routes.xcoms import _build_xcom_read from airflow.executors.workloads.base import BundleInfo from airflow.executors.workloads.task import ExecuteTask from airflow.models.dag_version import DagVersion @@ -37,7 +39,11 @@ from airflow.sdk.definitions._internal.loop import create_loop from airflow.utils.log.task_log_address import prepare_task_log_contexts -from tests_common.test_utils.asserts import assert_queries_count, capture_orm_selects +from tests_common.test_utils.asserts import ( + assert_queries_count, + capture_orm_selects, + count_loaded_task_instances, +) pytestmark = pytest.mark.db_test @@ -141,6 +147,113 @@ def test_shared_loop_uses_pinned_graph_and_retained_producer(loop_coordinates, s ) +@pytest.fixture +def forked_loop_coordinates(loop_coordinates, session): + replacement = session.scalars( + select(DynamicRegion).where(DynamicRegion.forked_from_region_id.is_not(None)) + ).one() + replacement.resumes_from_index = 2 + session.flush() + return loop_coordinates + + +def test_outside_caller_reads_the_live_iterations_each_region_owns(forked_loop_coordinates, session): + dr, producer, consumer, previous, outside = forked_loop_coordinates + resolver = TaskCoordinateResolver(DBDagBag(), session) + arguments = dict( + dag_id=dr.dag_id, run_id=dr.run_id, task_id=producer.task_id, caller=outside, all_iterations=True + ) + + selected = resolver.resolve(**arguments) + selected_ids = set(session.scalars(resolver.select_producer_ids(**arguments))) + + assert selected == (previous,) + assert selected_ids == {previous.id} + + +@pytest.mark.parametrize("by", ["resolve", "select_producer_ids"]) +def test_outside_caller_is_refused_a_loop_task_unless_it_reads_every_iteration(loop_coordinates, session, by): + dr, producer, consumer, previous, outside = loop_coordinates + resolver = TaskCoordinateResolver(DBDagBag(), session) + + with pytest.raises(ValueError, match="explicit scope or a consumer inside the loop"): + getattr(resolver, by)(dag_id=dr.dag_id, run_id=dr.run_id, task_id=producer.task_id, caller=outside) + + +def test_reading_every_iteration_cannot_be_combined_with_the_previous_one(loop_coordinates, session): + dr, producer, consumer, previous, outside = loop_coordinates + resolver = TaskCoordinateResolver(DBDagBag(), session) + + with pytest.raises(ValueError, match="requires a shared loop scope"): + resolver.resolve( + dag_id=dr.dag_id, + run_id=dr.run_id, + task_id=producer.task_id, + caller=outside, + all_iterations=True, + previous_iteration=True, + ) + + +def test_all_iterations_does_not_widen_the_read_of_a_caller_inside_the_loop(loop_coordinates, session): + dr, producer, consumer, previous, outside = loop_coordinates + resolver = TaskCoordinateResolver(DBDagBag(), session) + + selected = resolver.resolve( + dag_id=dr.dag_id, run_id=dr.run_id, task_id=producer.task_id, caller=consumer, all_iterations=True + ) + + assert selected == (producer,) + + +def test_dependency_of_an_outside_task_on_a_loop_task_is_every_live_iteration(loop_coordinates, session): + dr, producer, consumer, previous, outside = loop_coordinates + replacement = session.scalars( + select(DynamicRegion).where(DynamicRegion.forked_from_region_id.is_not(None)) + ).one() + replacement.resumes_from_index = 3 + first = TaskInstance( + task=producer.task, + run_id=dr.run_id, + dag_version_id=producer.dag_version_id, + region_id=producer.region_id, + ) + first.region_index = 0 + session.add(first) + session.flush() + resolver = TaskCoordinateResolver(DBDagBag(), session) + + assert set(resolver.resolve_dependency(outside, producer.task_id)) == {first, previous, producer} + + +@pytest.mark.parametrize( + ("caller", "task_id", "run_id", "expected"), + [ + pytest.param("outside", "body.producer", None, True, id="outside caller of a loop task"), + pytest.param("consumer", "body.producer", None, False, id="caller inside the loop"), + pytest.param("outside", "outside", None, False, id="task outside any loop"), + pytest.param("outside", "removed", None, False, id="task missing from the pinned dag"), + pytest.param("outside", "body.producer", "other_run", False, id="caller of another run"), + pytest.param(None, "body.producer", None, False, id="no caller"), + ], +) +def test_loop_task_read_from_outside_needs_a_caller_of_the_same_run_outside_the_loop( + loop_coordinates, session, caller, task_id, run_id, expected +): + dr, producer, consumer, previous, outside = loop_coordinates + resolver = TaskCoordinateResolver(DBDagBag(), session) + + assert ( + resolver.is_loop_task_read_from_outside( + dag_id=dr.dag_id, + run_id=run_id or dr.run_id, + task_id=task_id, + caller={"consumer": consumer, "outside": outside, None: None}[caller], + ) + is expected + ) + + def test_loop_passes_load_the_region_ancestry_once_for_all_task_instances(loop_coordinates, session): _, producer, consumer, previous, outside = loop_coordinates resolver = TaskCoordinateResolver(DBDagBag(), session) @@ -152,6 +265,31 @@ def test_loop_passes_load_the_region_ancestry_once_for_all_task_instances(loop_c assert len(statements) == 1 +def test_dependency_rejects_duplicate_latest_gate_coordinates(loop_coordinates, session): + dr, _, consumer, _, outside = loop_coordinates + resolver = TaskCoordinateResolver(DBDagBag(), session) + gate = next(ti for ti in dr.task_instances if ti.task_id == "body.__loop_gate") + duplicate = TaskInstance( + task=resolver.get_task(dr.dag_id, dr.run_id, gate.task_id, dag_version_id=gate.dag_version_id), + run_id=dr.run_id, + dag_version_id=gate.dag_version_id, + region_id=consumer.region_id, + region_index=gate.region_index, + ) + session.add(duplicate) + session.flush() + + with pytest.raises(AmbiguousProducerError, match="Multiple live loop gates"): + resolver.resolve_dependency(outside, gate.task_id) + + +def test_dependency_on_task_missing_from_pinned_dag_falls_back_to_plain_resolution(loop_coordinates, session): + _, _, consumer, _, _ = loop_coordinates + resolver = TaskCoordinateResolver(DBDagBag(), session) + + assert resolver.resolve_dependency(consumer, "removed_upstream") == () + + def test_loop_without_context_requires_explicit_scope(loop_coordinates, session): dr, producer, _, _, outside = loop_coordinates resolver = TaskCoordinateResolver(DBDagBag(), session) @@ -274,3 +412,104 @@ def test_unversioned_run_keeps_task_definition_after_latest_graph_changes( region_id=mapped.region_id, region_index=0, ) == (mapped,) + + +@pytest.fixture +def mapped_run(dag_maker, session): + def create(mapped_count: int): + with dag_maker(serialized=True): + mapped = PythonOperator.partial(task_id="mapped", python_callable=str).expand( + op_args=[[i] for i in range(mapped_count)] + ) + mapped >> EmptyOperator(task_id="reduce") + dr = dag_maker.create_dagrun() + caller = session.scalars( + select(TaskInstance).where(TaskInstance.run_id == dr.run_id, TaskInstance.task_id == "reduce") + ).one() + session.expire_all() + return dr, caller + + return create + + +@pytest.mark.parametrize("mapped_count", [3, 60]) +def test_xcom_read_of_one_mapped_slot_loads_no_producer_rows(mapped_run, session, mapped_count): + dr, caller = mapped_run(mapped_count) + token = TIToken(id=caller.id, claims=TIClaims()) + + with count_loaded_task_instances("mapped") as loaded, assert_queries_count(8): + read = _build_xcom_read( + dag_id=dr.dag_id, + run_id=dr.run_id, + task_id="mapped", + key="return_value", + session=session, + dag_bag=DBDagBag(), + token=token, + map_index=2, + ) + session.scalars(read.statement).all() + + assert loaded == [] + + +def test_loaded_task_instance_collector_keeps_only_the_requested_task(mapped_run, session): + dr, _ = mapped_run(3) + + with ( + count_loaded_task_instances("mapped") as loaded, + count_loaded_task_instances("other") as other_loaded, + ): + rows = session.scalars( + select(TaskInstance).where(TaskInstance.run_id == dr.run_id, TaskInstance.task_id == "mapped") + ).all() + + assert len(rows) == 3 + assert sorted(ti.id for ti in loaded) == sorted(ti.id for ti in rows) + assert other_loaded == [] + + +@pytest.mark.parametrize("mapped_count", [3, 60]) +def test_loop_read_of_one_mapped_slot_loads_no_producer_rows(dag_maker, session, mapped_count): + @task_group + def body(): + PythonOperator.partial(task_id="mapped", python_callable=str).expand( + op_args=[[i] for i in range(mapped_count)] + ) >> EmptyOperator(task_id="reduce") + + with dag_maker(serialized=True): + create_loop(body, max_iterations=3) + dr = dag_maker.create_dagrun() + caller = session.scalars( + select(TaskInstance).where(TaskInstance.run_id == dr.run_id, TaskInstance.task_id == "body.reduce") + ).one() + session.expire_all() + token = TIToken(id=caller.id, claims=TIClaims()) + + with count_loaded_task_instances("body.mapped") as loaded: + read = _build_xcom_read( + dag_id=dr.dag_id, + run_id=dr.run_id, + task_id="body.mapped", + key="return_value", + session=session, + dag_bag=DBDagBag(), + token=token, + map_index=2, + ) + session.scalars(read.statement).all() + + assert loaded == [] + + +def test_selected_mapped_producers_match_resolved_producers(mapped_run, session): + dr, caller = mapped_run(6) + resolver = TaskCoordinateResolver(DBDagBag(), session) + + for map_indexes in (None, 4, range(1, 3), [0, 5]): + arguments = dict( + dag_id=dr.dag_id, run_id=dr.run_id, task_id="mapped", caller=caller, map_indexes=map_indexes + ) + resolved = {ti.id for ti in resolver.resolve(**arguments)} + assert set(session.scalars(resolver.select_producer_ids(**arguments))) == resolved + assert resolved diff --git a/airflow-core/tests/unit/models/test_xcom_arg.py b/airflow-core/tests/unit/models/test_xcom_arg.py index d2887c0d0271f..c7c3a3136a295 100644 --- a/airflow-core/tests/unit/models/test_xcom_arg.py +++ b/airflow-core/tests/unit/models/test_xcom_arg.py @@ -17,6 +17,7 @@ from __future__ import annotations import pytest +from sqlalchemy import delete from airflow.models.dynamic_region import DynamicRegion, ProducerContext from airflow.models.expandinput import NotFullyPopulated @@ -37,6 +38,7 @@ ) from airflow.utils.state import TaskInstanceState +from tests_common.test_utils.asserts import assert_queries_count, count_loaded_task_instances from tests_common.test_utils.db import clear_db_dags, clear_db_runs pytestmark = pytest.mark.db_test @@ -463,3 +465,46 @@ def group(value): ) is None ) + + +def test_mapped_producer_without_live_instances_has_zero_length(dag_maker, session): + with dag_maker(session=session, serialized=True) as dag: + + @dag.task + def source(value): + return value + + source.expand(value=[1, 2]) + + dr = dag_maker.create_dagrun() + session.execute(delete(TaskInstance).where(TaskInstance.run_id == dr.run_id)) + session.flush() + argument = SchedulerPlainXComArg(dag_maker.serialized_dag.get_task("source"), XCOM_RETURN_KEY) + + assert get_task_map_length(argument, dr.run_id, session=session) == 0 + + +@pytest.mark.parametrize("mapped_count", [3, 60]) +def test_mapped_producer_length_loads_no_producer_rows(dag_maker, session, mapped_count): + with dag_maker(session=session, serialized=True) as dag: + + @dag.task + def source(value): + return value + + source.expand(value=list(range(mapped_count))) + + dr = dag_maker.create_dagrun() + for ti in dr.task_instances: + ti.state = TaskInstanceState.SUCCESS + XComModel.set_for_attempt( + task_instance_id=ti.id, key=XCOM_RETURN_KEY, value=ti.map_index, session=session + ) + session.flush() + session.expire_all() + argument = SchedulerPlainXComArg(dag_maker.serialized_dag.get_task("source"), XCOM_RETURN_KEY) + with count_loaded_task_instances("source") as loaded, assert_queries_count(4): + length = get_task_map_length(argument, dr.run_id, session=session) + + assert length == mapped_count + assert loaded == [] diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 85af2ca9dfde6..d40f36ebf6d33 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -1861,6 +1861,30 @@ def converged(loop): assert restored_loop.dag is restored assert restored_loop.parent_group is restored.task_group + def test_task_using_a_loop_task_serializes_with_the_gate_as_its_upstream(self): + @task_decorator + def report(value): + return value + + @task_group(group_id="body") + def body(): + BaseOperator(task_id="first") >> BaseOperator(task_id="terminal") + + with DAG("loop_outside_edge", schedule=None) as dag: + loop = create_loop(body, max_iterations=3) + report(loop["first"].output) + + encoded = DagSerialization.to_dict(dag) + DagSerialization.validate_schema(encoded) + restored = DagSerialization.from_json(DagSerialization.to_json(dag)) + + assert restored.get_task("report").upstream_task_ids == {"body.__loop_gate"} + assert restored.get_task("body.__loop_gate").downstream_task_ids == {"report"} + assert restored.get_task("body.first").downstream_task_ids == {"body.terminal"} + assert [arg.operator.task_id for arg in dag.get_task("report").op_args] == ["body.first"] + assert "body.first" in str(restored.get_task("report").op_args[0]) + assert "body.__loop_gate" not in str(restored.get_task("report").op_args[0]) + def test_task_group_serialization(self): """ Test TaskGroup serialization/deserialization. diff --git a/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py index 4850dde94269d..53d33165d0e90 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_trigger_rule_dep.py @@ -56,6 +56,80 @@ FAILED = TaskInstanceState.FAILED +@pytest.mark.parametrize("latest_state", [None, SUCCESS, FAILED, SKIPPED, TaskInstanceState.RESTARTING]) +@pytest.mark.parametrize("rule", [TriggerRule.ALL_SUCCESS, TriggerRule.ONE_SUCCESS, TriggerRule.ALL_DONE]) +def test_outside_dependency_observes_only_latest_live_loop_gate(dag_maker, session, latest_state, rule): + @task_group + def body(): + PythonOperator(task_id="terminal", python_callable=list) + + with dag_maker(serialized=True): + loop = create_loop(body, max_iterations=3) + loop >> PythonOperator(task_id="outside", python_callable=list, trigger_rule=rule) + dr = dag_maker.create_dagrun() + previous = next(ti for ti in dr.task_instances if ti.task_id == loop.gate_task_id) + previous.state = SUCCESS + latest = TaskInstance( + task=dag_maker.serialized_dag.get_task(previous.task_id), + run_id=dr.run_id, + dag_version_id=previous.dag_version_id, + region_id=previous.region_id, + region_index=1, + state=latest_state, + ) + session.add(latest) + outside = next(ti for ti in dr.task_instances if ti.task_id == "outside") + outside.task = dag_maker.serialized_dag.get_task(outside.task_id) + session.flush() + context = DepContext() + + actual = TriggerRuleDep().is_met(ti=outside, dep_context=context, session=session) + + expected = latest_state == SUCCESS or (rule == TriggerRule.ALL_DONE and latest_state in {FAILED, SKIPPED}) + assert actual is expected + assert context.upstream_tis(outside, previous.task_id, session=session) == (latest,) + with pytest.raises(ValueError, match="loop producer requires"): + context.coordinate_resolver(outside, session=session).resolve( + dag_id=dr.dag_id, run_id=dr.run_id, task_id=previous.task_id, caller=outside + ) + + +GATE_RULE_MATRIX = { + TriggerRule.ALL_SUCCESS: {SUCCESS}, + TriggerRule.ONE_SUCCESS: {SUCCESS}, + TriggerRule.ONE_FAILED: {FAILED, UPSTREAM_FAILED}, + TriggerRule.NONE_FAILED: {SUCCESS, SKIPPED}, + TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS: {SUCCESS}, + TriggerRule.ALL_DONE: {SUCCESS, FAILED, UPSTREAM_FAILED, SKIPPED}, + TriggerRule.ALWAYS: {SUCCESS, FAILED, UPSTREAM_FAILED, SKIPPED, None}, +} + + +@pytest.mark.parametrize("gate_state", [None, SUCCESS, FAILED, UPSTREAM_FAILED, SKIPPED]) +@pytest.mark.parametrize("rule", list(GATE_RULE_MATRIX)) +def test_task_using_a_loop_task_follows_the_loop_gate_not_the_loop_task(dag_maker, session, rule, gate_state): + @task_group + def body(): + EmptyOperator(task_id="first") >> EmptyOperator(task_id="terminal") + + with dag_maker(serialized=True): + loop = create_loop(body, max_iterations=3) + loop["first"] >> PythonOperator(task_id="outside", python_callable=list, trigger_rule=rule) + dr = dag_maker.create_dagrun() + tis = {ti.task_id: ti for ti in dr.task_instances} + tis["body.first"].state = SUCCESS + tis["body.terminal"].state = SUCCESS if gate_state in (SUCCESS, FAILED, SKIPPED) else FAILED + tis[loop.gate_task_id].state = gate_state + outside = tis["outside"] + outside.task = dag_maker.serialized_dag.get_task("outside") + session.flush() + + met = TriggerRuleDep().is_met(ti=outside, dep_context=DepContext(), session=session) + + assert outside.task.upstream_task_ids == {loop.gate_task_id} + assert met is (gate_state in GATE_RULE_MATRIX[rule]) + + @pytest.mark.parametrize("current_state", [SUCCESS, FAILED]) def test_trigger_rule_ignores_other_loop_passes(dag_maker, session, current_state): @task_group diff --git a/devel-common/src/tests_common/test_utils/asserts.py b/devel-common/src/tests_common/test_utils/asserts.py index e938bb0abbf6f..ab79ad768a04b 100644 --- a/devel-common/src/tests_common/test_utils/asserts.py +++ b/devel-common/src/tests_common/test_utils/asserts.py @@ -41,6 +41,8 @@ from sqlalchemy.orm import ORMExecuteState + from airflow.models.taskinstance import TaskInstance + log = logging.getLogger(__name__) @@ -282,3 +284,21 @@ def capture(orm_execute_state: ORMExecuteState) -> None: for element in joined - set(linter.froms) ) assert not problems, "Cartesian product in generated SQL:\n" + "\n".join(sorted(set(problems))) + + +@contextmanager +def count_loaded_task_instances(task_id: str) -> Generator[list[TaskInstance], None, None]: + """Collect the ``task_id`` task instance rows the ORM loads inside the block.""" + from airflow.models.taskinstance import TaskInstance + + loaded: list[TaskInstance] = [] + + def listener(target: TaskInstance, _context: object) -> None: + if target.task_id == task_id: + loaded.append(target) + + event.listen(TaskInstance, "load", listener) + try: + yield loaded + finally: + event.remove(TaskInstance, "load", listener) diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 49a2f603a373b..d7d779b15cf5f 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1142,6 +1142,9 @@ type GetXCom struct { // MapIndex corresponds to the JSON schema field "map_index". MapIndex interface{} `msgpack:"map_index,omitempty"` + // PreviousIteration corresponds to the JSON schema field "previous_iteration". + PreviousIteration bool `msgpack:"previous_iteration,omitempty"` + // RunID corresponds to the JSON schema field "run_id". RunID string `msgpack:"run_id"` @@ -1160,6 +1163,9 @@ type GetXComCount struct { // Key corresponds to the JSON schema field "key". Key string `msgpack:"key"` + // PreviousIteration corresponds to the JSON schema field "previous_iteration". + PreviousIteration bool `msgpack:"previous_iteration,omitempty"` + // RunID corresponds to the JSON schema field "run_id". RunID string `msgpack:"run_id"` @@ -1180,6 +1186,9 @@ type GetXComSequenceItem struct { // Offset corresponds to the JSON schema field "offset". Offset int `msgpack:"offset"` + // PreviousIteration corresponds to the JSON schema field "previous_iteration". + PreviousIteration bool `msgpack:"previous_iteration,omitempty"` + // RunID corresponds to the JSON schema field "run_id". RunID string `msgpack:"run_id"` @@ -1200,6 +1209,9 @@ type GetXComSequenceSlice struct { // Key corresponds to the JSON schema field "key". Key string `msgpack:"key"` + // PreviousIteration corresponds to the JSON schema field "previous_iteration". + PreviousIteration bool `msgpack:"previous_iteration,omitempty"` + // RunID corresponds to the JSON schema field "run_id". RunID string `msgpack:"run_id"` @@ -1306,6 +1318,24 @@ type LiteralArgBinding struct { type LogicalDates []time.Time +// Pinned loop definition and enclosing iteration for a task execution. +type LoopContext struct { + // Index corresponds to the JSON schema field "index". + Index int `msgpack:"index"` + + // MaxIterations corresponds to the JSON schema field "max_iterations". + MaxIterations int `msgpack:"max_iterations"` + + // NodeID corresponds to the JSON schema field "node_id". + NodeID string `msgpack:"node_id"` + + // TerminalIsMapped corresponds to the JSON schema field "terminal_is_mapped". + TerminalIsMapped bool `msgpack:"terminal_is_mapped"` + + // TerminalTaskID corresponds to the JSON schema field "terminal_task_id". + TerminalTaskID string `msgpack:"terminal_task_id"` +} + // Add a new value to be redacted in task logs. type MaskSecret struct { // Name corresponds to the JSON schema field "name". @@ -1654,6 +1684,9 @@ type TIRunContext struct { // DagRun corresponds to the JSON schema field "dag_run". DagRun DagRun `msgpack:"dag_run"` + // Loop corresponds to the JSON schema field "loop". + Loop *LoopContext `msgpack:"loop,omitempty"` + // MaxTries corresponds to the JSON schema field "max_tries". MaxTries int `msgpack:"max_tries"` @@ -1847,6 +1880,19 @@ type TaskStatesResult struct { Type string `msgpack:"type,omitempty"` } +// Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key"` + + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` +} + +type Warnings []interface{} + +type TriggerKwargs map[string]interface{} + type TriggerDagRun struct { // Conf corresponds to the JSON schema field "conf". Conf *Conf `msgpack:"conf,omitempty"` @@ -1876,8 +1922,6 @@ type TriggerDagRun struct { Type string `msgpack:"type,omitempty"` } -type TriggerKwargs map[string]interface{} - type UpdateDagRunNote struct { // Note corresponds to the JSON schema field "note". Note interface{} `msgpack:"note"` @@ -1889,19 +1933,8 @@ type UpdateDagRunNote struct { Type string `msgpack:"type,omitempty"` } -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` - - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` -} - type VersionData map[string]interface{} -type Warnings []interface{} - // Update the response content part of an existing Human-in-the-loop response. type UpdateHITLDetail struct { // ChosenOptions corresponds to the JSON schema field "chosen_options". diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json index b0d5797b788c8..7e09a736de793 100644 --- a/go-sdk/schema/supervisor-schema.json +++ b/go-sdk/schema/supervisor-schema.json @@ -2836,6 +2836,11 @@ }, "GetXCom": { "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -2888,6 +2893,11 @@ "GetXComCount": { "description": "Get the number of (mapped) XCom values available.", "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -2922,6 +2932,11 @@ }, "GetXComSequenceItem": { "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -2961,6 +2976,11 @@ }, "GetXComSequenceSlice": { "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -3222,6 +3242,42 @@ "title": "LazyDeserializedDAG", "type": "object" }, + "LoopContext": { + "description": "Pinned loop definition and enclosing iteration for a task execution.", + "properties": { + "node_id": { + "title": "Node Id", + "type": "string" + }, + "index": { + "minimum": 0, + "title": "Index", + "type": "integer" + }, + "max_iterations": { + "exclusiveMinimum": 0, + "title": "Max Iterations", + "type": "integer" + }, + "terminal_task_id": { + "title": "Terminal Task Id", + "type": "string" + }, + "terminal_is_mapped": { + "title": "Terminal Is Mapped", + "type": "boolean" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped" + ], + "title": "LoopContext", + "type": "object" + }, "MaskSecret": { "description": "Add a new value to be redacted in task logs.", "properties": { @@ -5154,6 +5210,17 @@ "dag_run": { "$ref": "#/$defs/DagRun" }, + "loop": { + "anyOf": [ + { + "$ref": "#/$defs/LoopContext" + }, + { + "type": "null" + } + ], + "default": null + }, "task_reschedule_count": { "default": 0, "title": "Task Reschedule Count", diff --git a/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py b/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py index ea9245f96ec48..d69dc99b1c88f 100644 --- a/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py +++ b/providers/edge3/tests/unit/edge3/worker_api/routes/test_logs.py @@ -75,8 +75,9 @@ def body(): create_loop(body, max_iterations=3) dr = dag_maker.create_dagrun(run_id="loop_run") ti = next(ti for ti in dr.task_instances if ti.task_id == "body.work") - region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id="body") - session.add(region) + region = DynamicRegion.get_or_create( + dag_id=dr.dag_id, run_id=dr.run_id, node_id="body", session=session + ) session.flush() ti.region_id, ti.region_index = region.id, 2 session.add( @@ -118,19 +119,20 @@ def body(): create_loop(body, max_iterations=3) dr = dag_maker.create_dagrun(run_id="mapped_loop_run") first = next(ti for ti in dr.task_instances if ti.task_id == "body.mapped") - loop_region = DynamicRegion(dag_id=dr.dag_id, run_id=dr.run_id, node_id="body") - session.add(loop_region) + loop_region = DynamicRegion.get_or_create( + dag_id=dr.dag_id, run_id=dr.run_id, node_id="body", session=session + ) session.flush() tis = [] for iteration in range(2): - region = DynamicRegion( + region = DynamicRegion.get_or_create( dag_id=dr.dag_id, run_id=dr.run_id, node_id="body.mapped", parent_region_id=loop_region.id, parent_region_index=iteration, + session=session, ) - session.add(region) session.flush() ti = ( first diff --git a/scripts/ci/prek/check_template_context_variable_in_sync.py b/scripts/ci/prek/check_template_context_variable_in_sync.py index 57459a863f38c..88bc3cbed3ac1 100755 --- a/scripts/ci/prek/check_template_context_variable_in_sync.py +++ b/scripts/ci/prek/check_template_context_variable_in_sync.py @@ -51,6 +51,8 @@ # AIP-103: task_state_store/asset_state_store aren't documented in templates-ref yet. Will be done in a later PR. "task_state_store", "asset_state_store", + # Only available inside private loop groups. + "loop", } diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 872c11cca131a..3d6b7db7e1052 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -606,9 +606,14 @@ def head( run_id: str, task_id: str, key: str, + *, + previous_iteration: bool = False, ) -> XComCountResponse: """Get the number of mapped XCom values.""" - resp = self.client.head(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}") + params: dict[str, str | int] = {} + if previous_iteration: + params["previous_iteration"] = True + resp = self.client.head(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params) # content_range: str | None if not (content_range := resp.headers["Content-Range"]) or not content_range.startswith( @@ -625,9 +630,13 @@ def get( key: str, map_index: int | None = None, include_prior_dates: bool = False, + *, + previous_iteration: bool = False, ) -> XComResponse: """Get a XCom value from the API server.""" params: dict[str, str | int] = {} + if previous_iteration: + params["previous_iteration"] = True if map_index is not None and map_index >= 0: params.update({"map_index": map_index}) if include_prior_dates: @@ -704,9 +713,17 @@ def get_sequence_item( task_id: str, key: str, offset: int, + *, + previous_iteration: bool = False, ) -> XComSequenceIndexResponse | ErrorResponse: + params: dict[str, str | int] = {} + if previous_iteration: + params["previous_iteration"] = True try: - resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}") + resp = self.client.get( + f"xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}", + params=params, + ) except ServerResponseError as e: if e.response.status_code == HTTPStatus.NOT_FOUND: log.error( @@ -742,8 +759,12 @@ def get_sequence_slice( stop: int | None, step: int | None, include_prior_dates: bool = False, + *, + previous_iteration: bool = False, ) -> XComSequenceSliceResponse: params: dict[str, str | int] = {} + if previous_iteration: + params["previous_iteration"] = True if start is not None: params["start"] = start if stop is not None: diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index bde0a9560d7b9..d6e9e45f0d24c 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -233,6 +233,18 @@ class IntermediateTIState(str, Enum): AWAITING_INPUT = "awaiting_input" +class LoopContext(BaseModel): + """ + Pinned loop definition and enclosing iteration for a task execution. + """ + + node_id: Annotated[str, Field(title="Node Id")] + index: Annotated[int, Field(ge=0, title="Index")] + max_iterations: Annotated[int, Field(gt=0, title="Max Iterations")] + terminal_task_id: Annotated[str, Field(title="Terminal Task Id")] + terminal_is_mapped: Annotated[bool, Field(title="Terminal Is Mapped")] + + class PrevSuccessfulDagRunResponse(BaseModel): """ Schema for response with previous successful DagRun information for Task Template Context. @@ -842,6 +854,7 @@ class TIRunContext(BaseModel): """ dag_run: DagRun + loop: LoopContext | None = None task_reschedule_count: Annotated[int | None, Field(title="Task Reschedule Count")] = 0 max_tries: Annotated[int, Field(title="Max Tries")] variables: Annotated[list[VariableResponse] | None, Field(title="Variables")] = None diff --git a/task-sdk/src/airflow/sdk/bases/decorator.py b/task-sdk/src/airflow/sdk/bases/decorator.py index baaee4a9894f4..896bedb3abc37 100644 --- a/task-sdk/src/airflow/sdk/bases/decorator.py +++ b/task-sdk/src/airflow/sdk/bases/decorator.py @@ -45,7 +45,7 @@ ListOfDictsExpandInput, is_mappable, ) -from airflow.sdk.definitions._internal.loop import check_dag_result_outside_loop +from airflow.sdk.definitions._internal.loop import LoopTaskGroup, check_dag_result_outside_loop from airflow.sdk.definitions._internal.types import NOTSET from airflow.sdk.definitions.asset import Asset from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS @@ -68,6 +68,14 @@ from airflow.sdk.definitions.taskgroup import TaskGroup +def _context_keys_for_task_group(task_group: TaskGroup | None) -> set[str]: + while task_group is not None: + if isinstance(task_group, LoopTaskGroup): + return KNOWN_CONTEXT_KEYS | {"loop"} + task_group = task_group.parent_group + return KNOWN_CONTEXT_KEYS + + class ExpandableFactory(Protocol): """ Protocol providing inspection against wrapped function. @@ -332,19 +340,24 @@ def __init__( # since values for those will be provided when the task is run. Since # we're not actually running the function, None is good enough here. signature = inspect.signature(python_callable) + task_group = kwargs.get("task_group") or TaskGroupContext.get_current(kwargs.get("dag")) + context_keys = _context_keys_for_task_group(task_group) + if "loop" in context_keys and (parameter := signature.parameters.get("loop")) is not None: + if parameter.kind != inspect.Parameter.KEYWORD_ONLY: + raise ValueError("The loop context parameter must be keyword-only") # Don't allow context argument defaults other than None to avoid ambiguities. faulty_parameters = [ param.name for param in signature.parameters.values() - if param.name in KNOWN_CONTEXT_KEYS and param.default not in (None, inspect.Parameter.empty) + if param.name in context_keys and param.default not in (None, inspect.Parameter.empty) ] if faulty_parameters: message = f"Context key parameter {faulty_parameters[0]} can't have a default other than None" raise ValueError(message) parameters = [ - param.replace(default=None) if param.name in KNOWN_CONTEXT_KEYS else param + param.replace(default=None) if param.name in context_keys else param for param in signature.parameters.values() ] @@ -368,7 +381,7 @@ def __init__( and param.default == inspect.Parameter.empty ): new_parameters.append(param.replace(default=None)) - if param.name not in KNOWN_CONTEXT_KEYS: + if param.name not in context_keys: injected_for_ordering.add(param.name) else: new_parameters.append(param) @@ -576,7 +589,8 @@ def __call__(self, *args: FParams.args, **kwargs: FParams.kwargs) -> XComArg: def _validate_arg_names(self, func: ValidationSource, kwargs: dict[str, Any]): # Ensure that context variables are not shadowed. - context_keys_being_mapped = KNOWN_CONTEXT_KEYS.intersection(kwargs) + task_group = self.kwargs.get("task_group") or TaskGroupContext.get_current(self.kwargs.get("dag")) + context_keys_being_mapped = _context_keys_for_task_group(task_group).intersection(kwargs) if len(context_keys_being_mapped) == 1: (name,) = context_keys_being_mapped raise ValueError(f"cannot call {func}() on task context variable {name!r}") diff --git a/task-sdk/src/airflow/sdk/bases/xcom.py b/task-sdk/src/airflow/sdk/bases/xcom.py index 50e54c458b087..1cbe08828459d 100644 --- a/task-sdk/src/airflow/sdk/bases/xcom.py +++ b/task-sdk/src/airflow/sdk/bases/xcom.py @@ -291,6 +291,7 @@ def get_one( run_id: str, map_index: int | None = None, include_prior_dates: bool = False, + _previous_iteration: bool = False, ) -> Any | None: """ Retrieve an XCom value, optionally meeting certain criteria. @@ -325,6 +326,7 @@ def get_one( run_id=run_id, map_index=map_index, include_prior_dates=include_prior_dates, + previous_iteration=_previous_iteration, ), ) diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/loop.py b/task-sdk/src/airflow/sdk/definitions/_internal/loop.py index da5f6bca753e5..d4c3767ccaae9 100644 --- a/task-sdk/src/airflow/sdk/definitions/_internal/loop.py +++ b/task-sdk/src/airflow/sdk/definitions/_internal/loop.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import inspect from collections.abc import Callable from typing import TYPE_CHECKING, Any @@ -25,8 +26,12 @@ from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup if TYPE_CHECKING: + from airflow.sdk.definitions.context import Context from airflow.sdk.definitions.decorators.task_group import _TaskGroupFactory +LOOP_XCOM_PREFIX = "_airflow_loop_" +LOOP_DECISION_KEY = f"{LOOP_XCOM_PREFIX}decision" + def _validate_max_iterations(instance, attribute, value: int) -> None: if isinstance(value, bool) or not isinstance(value, int) or value <= 0: @@ -53,6 +58,55 @@ def __attrs_post_init__(self): super().__attrs_post_init__() +def get_enclosing_loop(operator: Any) -> LoopTaskGroup | None: + group = operator.task_group + while isinstance(group, TaskGroup): + if isinstance(group, LoopTaskGroup): + return group + group = group.parent_group + return None + + +def is_loop_task_read_from_outside(producer: Any, consumer: Any) -> bool: + """Tell whether ``consumer`` reads ``producer``, a task of a loop, without being in that loop itself.""" + loop = get_enclosing_loop(producer) + if loop is None: + return False + consumer_loop = get_enclosing_loop(consumer) + return consumer_loop is None or consumer_loop.group_id != loop.group_id + + +def get_edge_source(source: Any, target: Any) -> Any: + """ + Return the task that ``target`` depends on when it is made downstream of ``source``. + + A task outside a loop waits for the loop's gate rather than for one of its tasks: the gate is what + creates the next iteration, so a member alone finishes before the loop has produced its later + iterations and a consumer would see only some of them. + """ + loop = get_enclosing_loop(source) + if loop is None or (gate_task_id := getattr(loop, "gate_task_id", None)) is None: + return source + target_loop = get_enclosing_loop(target) + if target_loop is not None and target_loop.group_id == loop.group_id: + return source + return source.dag.get_task(gate_task_id) + + +def redirect_outside_edges_to_gate(group: LoopTaskGroup, gate: BaseOperator) -> None: + """Make tasks outside the loop that the body already pointed at depend on the gate instead.""" + for task in list(group): + for node_id in list(task.downstream_task_ids): + downstream = gate.dag.get_task(node_id) + inside = get_enclosing_loop(downstream) + if inside is not None and inside.group_id == group.group_id: + continue + task.downstream_task_ids.discard(node_id) + downstream.upstream_task_ids.discard(task.node_id) + gate.downstream_task_ids.add(node_id) + downstream.upstream_task_ids.add(gate.node_id) + + def check_dag_result_outside_loop(operator: Any) -> None: """Raise if the operator sits in a loop group, whose tasks cannot be the Dag result.""" group = operator.task_group @@ -106,6 +160,12 @@ def __init__(self, *, until: Callable[..., bool] | None = None, **kwargs: Any): super().__init__(**kwargs) self.until = until + def execute(self, context: Context) -> None: + # execution_time.loop imports LOOP_DECISION_KEY from this module. + from airflow.sdk.execution_time.loop import execute_loop_gate + + execute_loop_gate(self, context) + def create_loop( factory: _TaskGroupFactory, @@ -119,6 +179,10 @@ def create_loop( if until is not None and not callable(until): raise TypeError("until must be callable") + if until is not None and ( + inspect.iscoroutinefunction(until) or inspect.iscoroutinefunction(type(until).__call__) + ): + raise TypeError("until must be synchronous") group = LoopTaskGroup( add_suffix_on_collision=True, **factory.tg_kwargs, @@ -142,5 +206,6 @@ def create_loop( terminal >> gate group.terminal_task_id = terminal.task_id group.gate_task_id = gate.task_id + redirect_outside_edges_to_gate(group, gate) factory._task_group_created = True return group diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/node.py b/task-sdk/src/airflow/sdk/definitions/_internal/node.py index 803ca837825fc..ac0a61d62772c 100644 --- a/task-sdk/src/airflow/sdk/definitions/_internal/node.py +++ b/task-sdk/src/airflow/sdk/definitions/_internal/node.py @@ -92,6 +92,7 @@ def _set_relatives( ) -> None: """Set relatives for the task or task list.""" from airflow.sdk.bases.operator import BaseOperator + from airflow.sdk.definitions._internal.loop import get_edge_source from airflow.sdk.definitions.mappedoperator import MappedOperator if not isinstance(task_or_task_list, Sequence): @@ -131,15 +132,17 @@ def _set_relatives( # If the other task does not yet have a Dag, add it to the same Dag as this task and dag.add_task(task) # type: ignore[arg-type] if upstream: - task.downstream_task_ids.add(self.node_id) - self.upstream_task_ids.add(task.node_id) + source = get_edge_source(task, self) + source.downstream_task_ids.add(self.node_id) + self.upstream_task_ids.add(source.node_id) if edge_modifier: - edge_modifier.add_edge_info(dag, task.node_id, self.node_id) + edge_modifier.add_edge_info(dag, source.node_id, self.node_id) else: - self.downstream_task_ids.add(task.node_id) - task.upstream_task_ids.add(self.node_id) + source = get_edge_source(self, task) + source.downstream_task_ids.add(task.node_id) + task.upstream_task_ids.add(source.node_id) if edge_modifier: - edge_modifier.add_edge_info(dag, self.node_id, task.node_id) + edge_modifier.add_edge_info(dag, source.node_id, task.node_id) def set_downstream( self, diff --git a/task-sdk/src/airflow/sdk/definitions/context.py b/task-sdk/src/airflow/sdk/definitions/context.py index 98c8051a24761..6f85470cc6819 100644 --- a/task-sdk/src/airflow/sdk/definitions/context.py +++ b/task-sdk/src/airflow/sdk/definitions/context.py @@ -33,6 +33,7 @@ InletEventsAccessors, TaskStateStoreAccessor, ) + from airflow.sdk.execution_time.loop import LoopContextAccessor from airflow.sdk.types import ( DagRunProtocol, Operator, @@ -57,6 +58,7 @@ class Context(TypedDict, total=False): inlets: list inlet_events: InletEventsAccessors logical_date: DateTime + loop: LoopContextAccessor macros: Any map_index_template: NotRequired[str | None] outlets: list @@ -90,7 +92,7 @@ class Context(TypedDict, total=False): var: Any -KNOWN_CONTEXT_KEYS: set[str] = set(Context.__annotations__.keys()) +KNOWN_CONTEXT_KEYS: set[str] = set(Context.__annotations__.keys()) - {"loop"} def context_merge(context: Context, *args: Any, **kwargs: Any) -> None: diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py index be581d32105bc..14ddc937dda26 100644 --- a/task-sdk/src/airflow/sdk/definitions/dag.py +++ b/task-sdk/src/airflow/sdk/definitions/dag.py @@ -1488,7 +1488,8 @@ def _run_task( from airflow.utils.session import create_session taskrun_result: TaskRunResult | None - log.info("[DAG TEST] starting task_id=%s map_index=%s", ti.task_id, ti.map_index) + map_index = ti.region_index if task.get_needs_expansion() else -1 + log.info("[DAG TEST] starting task_id=%s map_index=%s", ti.task_id, map_index) while True: try: log.info("[DAG TEST] running task %s", ti) @@ -1507,7 +1508,9 @@ def _run_task( dag_id=ti.dag_id, run_id=ti.run_id, try_number=ti.try_number, - map_index=ti.map_index, + map_index=map_index, + region_id=ti.region_id, + region_index=ti.region_index, dag_version_id=UUID(str(ti.dag_version_id)), ) @@ -1559,7 +1562,7 @@ def _run_task( break raise - log.info("[DAG TEST] end task task_id=%s map_index=%s", ti.task_id, ti.map_index) + log.info("[DAG TEST] end task task_id=%s map_index=%s", ti.task_id, map_index) return taskrun_result diff --git a/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py b/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py index 34943dd12a0e1..b0148278fb404 100644 --- a/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py +++ b/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py @@ -39,6 +39,7 @@ ListOfDictsExpandInput, MappedArgument, ) +from airflow.sdk.definitions._internal.loop import create_loop from airflow.sdk.definitions._internal.node import DAGNode from airflow.sdk.definitions.mappedoperator import ensure_xcomarg_return_value, prevent_duplicates from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup @@ -126,14 +127,33 @@ def _create_task_group(self, tg_factory: Callable[..., TaskGroup], *args: Any, * def override(self, **kwargs: Any) -> _TaskGroupFactory[FParams, FReturn]: # TODO: FIXME when mypy gets compatible with new attrs - return attr.evolve(self, tg_kwargs={**self.tg_kwargs, **kwargs}) # type: ignore[arg-type] + evolved = attr.evolve(self, tg_kwargs={**self.tg_kwargs, **kwargs}) # type: ignore[arg-type] + # The evolved copy owns the "never mapped" warning, so the intermediate factory must not repeat it. + self._task_group_created = True + return evolved + + def loop(self, *, max_iterations: int, until: Callable[..., bool] | None = None) -> TaskGroup: + """ + Repeat this group's tasks, creating each iteration when the previous gate continues. + + The body must have one terminal task definition, which may be mapped. + Use ``partial()`` to supply body arguments and ``override()`` to configure the group. + + :param max_iterations: Positive upper bound on the number of iterations. + :param until: Synchronous callable receiving task context keyword arguments. + True stops the loop. False at the iteration limit fails the gate. + Omit for a fixed-count loop that succeeds at its limit. + """ + return create_loop(self, max_iterations=max_iterations, until=until) def partial(self, **kwargs: Any) -> _TaskGroupFactory[FParams, FReturn]: self._validate_arg_names("partial", kwargs) prevent_duplicates(self.partial_kwargs, kwargs, fail_reason="duplicate partial") kwargs.update(self.partial_kwargs) # TODO: FIXME when mypy gets compatible with new attrs - return attr.evolve(self, partial_kwargs=kwargs) # type: ignore[arg-type] + evolved = attr.evolve(self, partial_kwargs=kwargs) # type: ignore[arg-type] + self._task_group_created = True + return evolved def expand(self, **kwargs: OperatorExpandArgument) -> DAGNode: if not kwargs: diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py index a38e47a466bcf..36b393c165f42 100644 --- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py +++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py @@ -28,6 +28,7 @@ from airflow.sdk import TriggerRule from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator +from airflow.sdk.definitions._internal.loop import is_loop_task_read_from_outside from airflow.sdk.definitions._internal.mixins import DependencyMixin, ResolveMixin from airflow.sdk.definitions._internal.setup_teardown import SetupTeardownContext from airflow.sdk.definitions._internal.types import NOTSET, is_arg_set @@ -333,7 +334,7 @@ def resolve(self, context: Mapping[str, Any]) -> Any: ti = context["ti"] task_id = self.operator.task_id - if self.operator.is_mapped: + if self.operator.is_mapped or is_loop_task_read_from_outside(self.operator, ti.task): return LazyXComSequence(xcom_arg=self, ti=ti) tg = self.operator.get_closest_mapped_task_group() if tg is None: diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 6472863fbd8a7..7c957027ba792 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -941,6 +941,7 @@ class SkipDownstreamTasks(TISkippedDownstreamTasksStatePayload): class GetXCom(BaseModel): + previous_iteration: bool = False key: str dag_id: str run_id: str @@ -953,6 +954,7 @@ class GetXCom(BaseModel): class GetXComCount(BaseModel): """Get the number of (mapped) XCom values available.""" + previous_iteration: bool = False key: str dag_id: str run_id: str @@ -961,6 +963,7 @@ class GetXComCount(BaseModel): class GetXComSequenceItem(BaseModel): + previous_iteration: bool = False key: str dag_id: str run_id: str @@ -970,6 +973,7 @@ class GetXComSequenceItem(BaseModel): class GetXComSequenceSlice(BaseModel): + previous_iteration: bool = False key: str dag_id: str run_id: str diff --git a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py index 4efb0b71368ca..4ce198d14f58b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py +++ b/task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py @@ -65,6 +65,7 @@ class LazyXComSequence(Sequence[T]): _len: int | None = attrs.field(init=False, default=None) _xcom_arg: PlainXComArg = attrs.field(alias="xcom_arg") _ti: RuntimeTaskInstance = attrs.field(alias="ti") + _previous_iteration: bool = attrs.field(default=False, alias="previous_iteration", kw_only=True) def __repr__(self) -> str: if self._len is not None: @@ -96,6 +97,7 @@ def __len__(self) -> int: msg = SUPERVISOR_COMMS.send( GetXComCount( + previous_iteration=self._previous_iteration, key=self._xcom_arg.key, dag_id=task.dag_id, run_id=self._ti.run_id, @@ -131,6 +133,7 @@ def __getitem__(self, key: int | slice) -> T | Sequence[T]: source = (xcom_arg := self._xcom_arg).operator msg = SUPERVISOR_COMMS.send( GetXComSequenceSlice( + previous_iteration=self._previous_iteration, key=xcom_arg.key, dag_id=source.dag_id, task_id=source.task_id, @@ -152,6 +155,7 @@ def __getitem__(self, key: int | slice) -> T | Sequence[T]: source = (xcom_arg := self._xcom_arg).operator msg = SUPERVISOR_COMMS.send( GetXComSequenceItem( + previous_iteration=self._previous_iteration, key=xcom_arg.key, dag_id=source.dag_id, task_id=source.task_id, diff --git a/task-sdk/src/airflow/sdk/execution_time/loop.py b/task-sdk/src/airflow/sdk/execution_time/loop.py new file mode 100644 index 0000000000000..a3b52282b1ff5 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/loop.py @@ -0,0 +1,104 @@ +# 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. +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import attrs + +from airflow.sdk.bases.decorator import determine_kwargs +from airflow.sdk.definitions._internal.loop import LOOP_DECISION_KEY +from airflow.sdk.definitions.xcom_arg import PlainXComArg +from airflow.sdk.exceptions import AirflowFailException +from airflow.sdk.execution_time.comms import SetXCom +from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence +from airflow.sdk.execution_time.xcom import XCom + +if TYPE_CHECKING: + from airflow.sdk.api.datamodels._generated import LoopContext + from airflow.sdk.definitions._internal.loop import LoopGateOperator + from airflow.sdk.definitions.context import Context + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + + +class LoopMaxIterationsExceeded(AirflowFailException): + """The loop condition remained false at its iteration cap.""" + + +@attrs.define(frozen=True) +class LoopContextAccessor: + _descriptor: LoopContext + _ti: RuntimeTaskInstance + + @property + def index(self) -> int: + return self._descriptor.index + + @property + def max_iterations(self) -> int: + return self._descriptor.max_iterations + + @property + def previous(self) -> Any: + return None if self.index == 0 else self._result(previous_iteration=True) + + @property + def result(self) -> Any: + return self._result(previous_iteration=False) + + def _result(self, *, previous_iteration: bool) -> Any: + task = self._ti.task.dag.get_task(self._descriptor.terminal_task_id) + if self._descriptor.terminal_is_mapped: + return LazyXComSequence( + xcom_arg=PlainXComArg(task), ti=self._ti, previous_iteration=previous_iteration + ) + return XCom.get_one( + key=XCom.XCOM_RETURN_KEY, + dag_id=self._ti.dag_id, + run_id=self._ti.run_id, + task_id=task.task_id, + _previous_iteration=previous_iteration, + ) + + +def execute_loop_gate(gate: LoopGateOperator, context: Context) -> None: + # task_runner imports LoopContextAccessor from this module. + from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS + + loop = context["loop"] + if gate.until is None: + stop = loop.index + 1 >= loop.max_iterations + else: + stop = gate.until(**determine_kwargs(gate.until, (), context)) + if not isinstance(stop, bool): + raise AirflowFailException(f"until must return a bool, got {type(stop).__name__}") + if not stop and loop.index + 1 >= loop.max_iterations: + raise LoopMaxIterationsExceeded( + f"Loop condition remained false at max_iterations={loop.max_iterations}" + ) + decision = "stop" if stop else "continue" + gate.log.info("Loop iteration %s: %s (max_iterations=%s)", loop.index, decision, loop.max_iterations) + ti = context["ti"] + SUPERVISOR_COMMS.send( + SetXCom( + dag_id=ti.dag_id, + run_id=ti.run_id, + task_id=ti.task_id, + key=LOOP_DECISION_KEY, + value=decision, + ) + ) diff --git a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py index 594a871f3a859..2a2d16833dffb 100644 --- a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py +++ b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py @@ -257,7 +257,13 @@ def handle_get_prev_successful_dag_run( def handle_get_xcom_count(client: Client, msg: GetXComCount) -> tuple[BaseModel | None, dict[str, bool]]: """Fetch XCom count metadata.""" - resp = client.xcoms.head(msg.dag_id, msg.run_id, msg.task_id, msg.key) + resp = client.xcoms.head( + msg.dag_id, + msg.run_id, + msg.task_id, + msg.key, + previous_iteration=msg.previous_iteration, + ) return resp, {} @@ -271,6 +277,7 @@ def handle_get_xcom_sequence_item( msg.task_id, msg.key, msg.offset, + previous_iteration=msg.previous_iteration, ) if isinstance(xcom, XComSequenceIndexResponse): return XComSequenceIndexResult.from_response(xcom), {} @@ -290,6 +297,7 @@ def handle_get_xcom_sequence_slice( msg.stop, msg.step, msg.include_prior_dates, + previous_iteration=msg.previous_iteration, ) if isinstance(xcoms, XComSequenceSliceResponse): return XComSequenceSliceResult.from_response(xcoms), {} @@ -305,6 +313,7 @@ def handle_get_xcom(client: Client, msg: GetXCom) -> tuple[BaseModel | None, dic msg.key, msg.map_index, msg.include_prior_dates, + previous_iteration=msg.previous_iteration, ) if isinstance(xcom, XComResponse): xcom_result = XComResult.from_xcom_response(xcom) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index b0d5797b788c8..7e09a736de793 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -2836,6 +2836,11 @@ }, "GetXCom": { "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -2888,6 +2893,11 @@ "GetXComCount": { "description": "Get the number of (mapped) XCom values available.", "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -2922,6 +2932,11 @@ }, "GetXComSequenceItem": { "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -2961,6 +2976,11 @@ }, "GetXComSequenceSlice": { "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -3222,6 +3242,42 @@ "title": "LazyDeserializedDAG", "type": "object" }, + "LoopContext": { + "description": "Pinned loop definition and enclosing iteration for a task execution.", + "properties": { + "node_id": { + "title": "Node Id", + "type": "string" + }, + "index": { + "minimum": 0, + "title": "Index", + "type": "integer" + }, + "max_iterations": { + "exclusiveMinimum": 0, + "title": "Max Iterations", + "type": "integer" + }, + "terminal_task_id": { + "title": "Terminal Task Id", + "type": "string" + }, + "terminal_is_mapped": { + "title": "Terminal Is Mapped", + "type": "boolean" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped" + ], + "title": "LoopContext", + "type": "object" + }, "MaskSecret": { "description": "Add a new value to be redacted in task logs.", "properties": { @@ -5154,6 +5210,17 @@ "dag_run": { "$ref": "#/$defs/DagRun" }, + "loop": { + "anyOf": [ + { + "$ref": "#/$defs/LoopContext" + }, + { + "type": "null" + } + ], + "default": null + }, "task_reschedule_count": { "default": 0, "title": "Task Reschedule Count", diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index bdbb825c2d126..647e5dc919317 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -40,6 +40,7 @@ def get_bundle() -> VersionBundle: from airflow.sdk.execution_time.schema.versions.v2026_10_30 import ( AddArgBindingsToSupervisorTIRunContext, AddDagDefinitionsToDagFileParsingResult, + AddLoopRuntime, AddRegionCoordinatesToSupervisorTaskInstance, AddRetryReasonToTaskState, ) @@ -49,6 +50,7 @@ def get_bundle() -> VersionBundle: Version( "2026-10-30", AddArgBindingsToSupervisorTIRunContext, + AddLoopRuntime, AddRetryReasonToTaskState, AddDagDefinitionsToDagFileParsingResult, AddRegionCoordinatesToSupervisorTaskInstance, diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py index e093ecd07ec00..baffc8c9d0f1b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py @@ -31,10 +31,27 @@ from airflow.dag_processing.processor import DagFileParsingResult # noqa: SDK002 from airflow.sdk.api.datamodels._generated import PreviousTIResponse, TaskInstance, TIRunContext from airflow.sdk.execution_time.comms import ( + GetXCom, + GetXComCount, + GetXComSequenceItem, + GetXComSequenceSlice, TaskState, ) +class AddLoopRuntime(VersionChange): + """Carry loop context and previous-iteration reads.""" + + description = __doc__ + instructions_to_migrate_to_previous_version = ( + schema(TIRunContext).field("loop").didnt_exist, + schema(GetXCom).field("previous_iteration").didnt_exist, + schema(GetXComCount).field("previous_iteration").didnt_exist, + schema(GetXComSequenceItem).field("previous_iteration").didnt_exist, + schema(GetXComSequenceSlice).field("previous_iteration").didnt_exist, + ) + + class AddRegionCoordinatesToSupervisorTaskInstance(VersionChange): """Carry task coordinates in supervisor startup and previous-instance messages.""" diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index aa6aa477c0cec..d4fbec7c25a3c 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -63,6 +63,7 @@ from airflow.sdk.configuration import conf from airflow.sdk.coordinators._dag_importer import find_claiming_importer from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager +from airflow.sdk.definitions._internal.loop import LOOP_XCOM_PREFIX, is_loop_task_read_from_outside from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set from airflow.sdk.definitions.asset import ( Asset, @@ -73,6 +74,7 @@ ) from airflow.sdk.definitions.mappedoperator import MappedOperator from airflow.sdk.definitions.param import process_params +from airflow.sdk.definitions.xcom_arg import PlainXComArg from airflow.sdk.exceptions import ( AirflowException, AirflowFailException, @@ -85,6 +87,7 @@ ErrorType, TaskAwaitingInput, TaskDeferred, + TaskNotFound, ) from airflow.sdk.execution_time.bundles import initialize_ti_bundle from airflow.sdk.execution_time.callback_runner import create_executable_runner @@ -150,6 +153,8 @@ _ErrorEmailNotifier, _LegacyEmailBackendNotifier, ) +from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence +from airflow.sdk.execution_time.loop import LoopContextAccessor from airflow.sdk.execution_time.sentry import Sentry from airflow.sdk.execution_time.tracing import detail_span from airflow.sdk.execution_time.xcom import XCom @@ -350,6 +355,8 @@ def get_template_context(self) -> Context: if TYPE_CHECKING: assert self._cached_template_context is not None if from_server: + if from_server.loop is not None: + self._cached_template_context["loop"] = LoopContextAccessor(from_server.loop, self) dag_run = from_server.dag_run context_from_server: Context = { # TODO: Assess if we need to pass these through timezone.coerce_datetime @@ -493,6 +500,20 @@ def _normalize_xcom_pull_params( map_indexes_iterable, ) + def _get_loop_task_sequence( + self, task_id: str, *, key: str, dag_id: str, run_id: str, include_prior_dates: bool + ) -> LazyXComSequence | None: + """Return every iteration of a loop task as a sequence when this task instance is outside that loop.""" + if include_prior_dates or (dag_id, run_id) != (self.dag_id, self.run_id): + return None + try: + producer = self.task.dag.get_task(task_id) + except TaskNotFound: + return None + if not is_loop_task_read_from_outside(producer, self.task): + return None + return LazyXComSequence(xcom_arg=PlainXComArg(producer, key), ti=self) + def xcom_pull( self, task_ids: str | Iterable[str] | None = None, @@ -551,6 +572,16 @@ def xcom_pull( xcoms: list[Any] = [] if not is_arg_set(map_indexes_iterable): + if single_task_requested: + sequence = self._get_loop_task_sequence( + task_ids[0], + key=key, + dag_id=dag_id, + run_id=run_id, + include_prior_dates=include_prior_dates, + ) + if sequence is not None: + return sequence # map_indexes was not specified — fetch all map indexes for each task for t_id in task_ids: values = XCom.get_all( @@ -600,6 +631,17 @@ async def axcom_pull( xcoms: list[Any] = [] if not is_arg_set(map_indexes_iterable): + if single_task_requested and ( + self._get_loop_task_sequence( + task_ids[0], + key=key, + dag_id=dag_id, + run_id=run_id, + include_prior_dates=include_prior_dates, + ) + is not None + ): + return await XCom.aget_all(run_id=run_id, key=key, task_id=task_ids[0], dag_id=dag_id) or [] # map_indexes was not specified — fetch all map indexes for each task for t_id in task_ids: values = await XCom.aget_all( @@ -1575,7 +1617,8 @@ def _on_term(signum, frame): if ti._ti_context_from_server and (keys_to_delete := ti._ti_context_from_server.xcom_keys_to_clear): for x in keys_to_delete: log.debug("Clearing XCom with key", key=x) - XCom.delete( + backend = BaseXCom if x.startswith(LOOP_XCOM_PREFIX) else XCom + backend.delete( key=x, dag_id=ti.dag_id, task_id=ti.task_id, diff --git a/task-sdk/tests/task_sdk/definitions/_internal/test_loop.py b/task-sdk/tests/task_sdk/definitions/_internal/test_loop.py index 28bb2eda92e76..d2b8db90654e4 100644 --- a/task-sdk/tests/task_sdk/definitions/_internal/test_loop.py +++ b/task-sdk/tests/task_sdk/definitions/_internal/test_loop.py @@ -20,10 +20,68 @@ import pytest -from airflow.sdk import DAG, BaseOperator, TaskGroup, TriggerRule, task, task_group +from airflow.sdk import ( + DAG, + BaseOperator, + Label, + TaskGroup, + TriggerRule, + chain, + cross_downstream, + task, + task_group, +) from airflow.sdk.definitions._internal.loop import LoopGateOperator, create_loop +@pytest.mark.parametrize("mapped", [False, True]) +def test_loop_body_accepts_omitted_keyword_only_loop_context(mapped): + @task + def terminal(value, *, loop): + return value + loop.index + + @task_group + def body(): + with TaskGroup("nested"): + if mapped: + terminal.expand(value=[1, 2]) + else: + terminal(1) + + with DAG("loop_context_argument", schedule=None) as dag: + group = create_loop(body, max_iterations=2) + + assert group.terminal_task_id == "body.nested.terminal" + assert dag.get_task(group.gate_task_id).upstream_task_ids == {group.terminal_task_id} + + +def test_ordinary_task_loop_argument_remains_required(): + @task + def terminal(*, loop): + return loop + + with DAG("ordinary_loop_argument", schedule=None): + with pytest.raises(TypeError, match="missing a required (keyword-only )?argument: 'loop'"): + terminal() + result = terminal(loop="user value") + + assert result.operator.op_kwargs == {"loop": "user value"} + + +def test_loop_context_cannot_be_replaced_by_mapped_input(): + @task + def terminal(*, loop): + return loop.index + + @task_group + def body(): + terminal.expand(loop=["user value"]) + + with DAG("mapped_loop_context", schedule=None): + with pytest.raises(ValueError, match="task context variable 'loop'"): + create_loop(body, max_iterations=2) + + def test_loop_dependencies_follow_terminal_instead_of_body_return(): @task_group def body(): @@ -144,6 +202,25 @@ def body(): create_loop(body, max_iterations=2, until=until) +@pytest.mark.parametrize("kind", ["function", "partial", "object"]) +def test_async_condition_is_rejected(kind): + @task_group + def body(): + BaseOperator(task_id="terminal") + + async def until(*, loop): + return True + + class AsyncCondition: + async def __call__(self, *, loop): + return True + + condition = {"function": until, "partial": partial(until), "object": AsyncCondition()}[kind] + with DAG("async_condition", schedule=None): + with pytest.raises(TypeError, match="synchronous"): + create_loop(body, max_iterations=2, until=condition) + + @pytest.mark.parametrize("terminal_count", [0, 2]) def test_loop_requires_one_terminal_task_definition(terminal_count): @task_group @@ -200,3 +277,155 @@ def mapped(value): with DAG("mapped_loop", schedule=None): with pytest.raises(NotImplementedError, match="mapped task group"): mapped.expand(value=[1, 2]) + + +def make_edge_dag(): + @task_group + def body(): + BaseOperator(task_id="first") >> BaseOperator(task_id="second") + + with DAG("loop_edges", schedule=None) as dag: + loop = create_loop(body, max_iterations=3) + BaseOperator(task_id="before") + BaseOperator(task_id="after") + return dag, loop + + +def link_with_rshift(member, consumer): + member >> consumer + + +def link_with_lshift(member, consumer): + consumer << member + + +def link_with_set_downstream(member, consumer): + member.set_downstream(consumer) + + +def link_with_set_upstream(member, consumer): + consumer.set_upstream(member) + + +def link_with_list(member, consumer): + [member] >> consumer + + +def link_with_label(member, consumer): + member >> Label("loop result") >> consumer + + +def link_with_chain(member, consumer): + chain(member, consumer) + + +def link_with_cross_downstream(member, consumer): + cross_downstream([member], [consumer]) + + +def link_with_xcom_arg(member, consumer): + consumer.set_upstream(member.output) + + +@pytest.mark.parametrize( + "link", + [ + link_with_rshift, + link_with_lshift, + link_with_set_downstream, + link_with_set_upstream, + link_with_list, + link_with_label, + link_with_chain, + link_with_cross_downstream, + link_with_xcom_arg, + ], +) +def test_outside_task_depends_on_the_gate_not_on_the_loop_task(link): + dag, loop = make_edge_dag() + + link(dag.get_task("body.first"), dag.get_task("after")) + + assert dag.get_task("after").upstream_task_ids == {loop.gate_task_id} + assert "after" in dag.get_task(loop.gate_task_id).downstream_task_ids + assert dag.get_task("body.first").downstream_task_ids == {"body.second"} + + +def test_task_group_edge_to_an_outside_task_uses_the_gate(): + dag, loop = make_edge_dag() + + loop >> dag.get_task("after") + + assert dag.get_task("after").upstream_task_ids == {loop.gate_task_id} + + +def test_edges_that_stay_inside_or_enter_the_loop_are_unchanged(): + dag, loop = make_edge_dag() + + with dag, loop: + extra = BaseOperator(task_id="extra") + dag.get_task("body.first") >> extra + dag.get_task("before") >> dag.get_task("body.first") + dag.get_task("before") >> loop + + assert extra.upstream_task_ids == {"body.first"} + assert dag.get_task("body.first").upstream_task_ids == {"before"} + assert dag.get_task("body.second").upstream_task_ids == {"body.first"} + assert dag.get_task(loop.gate_task_id).upstream_task_ids == {"body.second"} + assert dag.get_task("before").downstream_task_ids == {"body.first"} + + +def test_xcom_argument_of_a_loop_task_keeps_reading_the_task_but_waits_for_the_gate(): + @task + def produce(): + return 1 + + @task + def report(value): + return value + + @task_group + def body(): + produce() + + with DAG("loop_xcom_edge", schedule=None): + loop = create_loop(body, max_iterations=2) + consumer = report(loop["produce"].output) + + assert [arg.operator.task_id for arg in consumer.operator.op_args] == ["body.produce"] + assert consumer.operator.upstream_task_ids == {loop.gate_task_id} + + +def test_mapped_consumer_expanded_over_a_loop_task_waits_for_the_gate(): + @task + def produce(): + return [1, 2] + + @task + def use(value): + return value + + @task_group + def body(): + produce() + + with DAG("loop_expand_edge", schedule=None): + loop = create_loop(body, max_iterations=2) + consumer = use.expand(value=loop["produce"].output) + + assert consumer.operator.upstream_task_ids == {loop.gate_task_id} + + +def test_edge_to_an_outside_task_made_while_the_body_is_built_waits_for_the_gate(): + with DAG("loop_early_edge", schedule=None) as dag: + outside = BaseOperator(task_id="outside") + + @task_group + def body(): + BaseOperator(task_id="first") >> outside + + loop = create_loop(body, max_iterations=2) + + assert outside.upstream_task_ids == {loop.gate_task_id} + assert dag.get_task("body.first").downstream_task_ids == {loop.gate_task_id} + assert outside.task_id in dag.get_task(loop.gate_task_id).downstream_task_ids diff --git a/task-sdk/tests/task_sdk/definitions/decorators/test_task_group.py b/task-sdk/tests/task_sdk/definitions/decorators/test_task_group.py index 7a777b1ea1988..a3519946f8474 100644 --- a/task-sdk/tests/task_sdk/definitions/decorators/test_task_group.py +++ b/task-sdk/tests/task_sdk/definitions/decorators/test_task_group.py @@ -17,13 +17,15 @@ # under the License. from __future__ import annotations +import gc +import warnings from datetime import timedelta import pendulum import pytest from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.sdk import dag, task, task_group +from airflow.sdk import DAG, dag, task, task_group from airflow.sdk.definitions._internal.expandinput import ( DictOfListsExpandInput, ListOfDictsExpandInput, @@ -32,6 +34,89 @@ from airflow.sdk.definitions.taskgroup import MappedTaskGroup +@pytest.mark.parametrize("mapped", [False, True]) +@pytest.mark.parametrize("conditional", [False, True]) +def test_loop_preserves_group_configuration_and_body_arguments(mapped, conditional): + @task + def terminal(value, *, loop): + return value + loop.index + + @task_group + def body(value): + """Refine the input.""" + if mapped: + return terminal.expand(value=[value, value + 1]) + return terminal(value) + + def converged(*, loop): + return loop.index == 1 + + with DAG("public_loop", schedule=None) as pipeline: + start = EmptyOperator(task_id="start") + finish = EmptyOperator(task_id="finish") + group = ( + body.override(group_id="refine") + .partial(value=2) + .loop(max_iterations=3, until=converged if conditional else None) + ) + start >> group >> finish + + gate = pipeline.get_task("refine.converged" if conditional else "refine.__loop_gate") + assert group.doc_md == "Refine the input." + assert gate.until is (converged if conditional else None) + assert gate.upstream_task_ids == {"refine.terminal"} + assert gate.downstream_task_ids == {"finish"} + result = pipeline.get_task("refine.terminal") + assert result.upstream_task_ids == {"start"} + if mapped: + assert result.op_kwargs_expand_input.value == {"value": [2, 3]} + else: + assert result.op_args == (2,) + + +@pytest.mark.parametrize("max_iterations", [0, True]) +def test_loop_rejects_invalid_iteration_limit(max_iterations): + @task_group + def body(): + EmptyOperator(task_id="terminal") + + with DAG("invalid_public_loop", schedule=None): + with pytest.raises(ValueError, match="positive integer"): + body.loop(max_iterations=max_iterations) + + +def test_loop_rejects_multiple_terminal_definitions(): + @task_group + def body(): + EmptyOperator(task_id="first") + EmptyOperator(task_id="second") + + with DAG("ambiguous_public_loop", schedule=None): + with pytest.raises(ValueError, match="exactly one terminal"): + body.loop(max_iterations=2) + + +@pytest.mark.parametrize("partial_first", [False, True]) +def test_loop_after_partial_and_override_in_either_order_does_not_warn(partial_first): + @task_group + def body(value): + EmptyOperator(task_id="terminal") + + with DAG("ordered_public_loop", schedule=None) as pipeline: + with warnings.catch_warnings(record=True) as recorded: + warnings.simplefilter("always") + if partial_first: + factory = body.partial(value=2).override(group_id="refine") + else: + factory = body.override(group_id="refine").partial(value=2) + factory.loop(max_iterations=2) + del factory + gc.collect() + + assert pipeline.get_task("refine.terminal") + assert [str(warning.message) for warning in recorded if "never mapped" in str(warning.message)] == [] + + def test_task_group_with_overridden_kwargs(): @task_group( default_args={ diff --git a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py index 2d1f8b2a4fa52..53b1139dacb41 100644 --- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py +++ b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py @@ -23,12 +23,19 @@ import pytest import structlog -from airflow.sdk import TaskInstanceState +from airflow.sdk import BaseOperator, TaskInstanceState, task_group from airflow.sdk.bases.xcom import BaseXCom +from airflow.sdk.definitions._internal.loop import create_loop from airflow.sdk.definitions.dag import DAG from airflow.sdk.definitions.xcom_arg import PlainXComArg from airflow.sdk.exceptions import AirflowSkipException -from airflow.sdk.execution_time.comms import GetXCom, XComResult, XComSequenceSliceResult +from airflow.sdk.execution_time.comms import ( + GetXCom, + GetXComCount, + XComCountResponse, + XComResult, + XComSequenceSliceResult, +) from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence from airflow.sdk.serde import deserialize, serialize @@ -416,3 +423,43 @@ def test_resolve_uses_xcom_pull_for_specific_index(self): assert resolved == "value-0" ti.xcom_pull.assert_called_once() assert ti.xcom_pull.call_args.kwargs["map_indexes"] == 0 + + +class TestPlainXComArgResolveLoopTask: + @pytest.fixture + def dag(self): + @task_group + def body(): + BaseOperator(task_id="first") >> BaseOperator(task_id="second") + + with DAG("loop_dag", schedule=None) as dag: + create_loop(body, max_iterations=3) >> BaseOperator(task_id="after") + return dag + + def test_resolve_from_outside_the_loop_returns_a_sequence_of_every_iteration( + self, dag, create_runtime_ti, mock_supervisor_comms + ): + ti = create_runtime_ti(task=dag.get_task("after")) + mock_supervisor_comms.send.return_value = XComCountResponse(len=3) + + resolved = PlainXComArg(dag.get_task("body.first")).resolve({"ti": ti}) + + assert isinstance(resolved, LazyXComSequence) + mock_supervisor_comms.send.assert_not_called() + assert len(resolved) == 3 + sent = mock_supervisor_comms.send.call_args.args[0] + assert sent == GetXComCount( + key="return_value", dag_id=ti.dag_id, run_id=ti.run_id, task_id="body.first" + ) + assert set(sent.model_dump()) == {"type", "previous_iteration", "key", "dag_id", "run_id", "task_id"} + + def test_resolve_from_inside_the_loop_pulls_the_value_of_the_same_iteration( + self, dag, create_runtime_ti, mock_supervisor_comms + ): + ti = create_runtime_ti(task=dag.get_task("body.second")) + mock_supervisor_comms.send.return_value = XComResult(key="return_value", value="value") + + resolved = PlainXComArg(dag.get_task("body.first")).resolve({"ti": ti}) + + assert resolved == "value" + assert isinstance(mock_supervisor_comms.send.call_args.args[0], GetXCom) diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index da3f120f06314..a5ac1de22d286 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -50,8 +50,12 @@ ) from airflow.sdk import TaskInstanceState -from airflow.sdk.api.datamodels._generated import PreviousTIResponse +from airflow.sdk.api.datamodels._generated import LoopContext, PreviousTIResponse from airflow.sdk.execution_time.comms import ( + GetXCom, + GetXComCount, + GetXComSequenceItem, + GetXComSequenceSlice, PreviousTIResult, TaskState, ) @@ -94,6 +98,23 @@ def test_previous_ti_response_coordinates_follow_supervisor_version(version): assert result["region_index"] == 3 +@pytest.mark.parametrize( + ("model", "extra", "field"), + [ + (GetXCom, {}, "previous_iteration"), + (GetXComCount, {}, "previous_iteration"), + (GetXComSequenceItem, {"offset": 0}, "previous_iteration"), + (GetXComSequenceSlice, {"start": None, "stop": None, "step": None}, "previous_iteration"), + ], +) +def test_loop_selectors_follow_supervisor_version(model, extra, field): + message = model(dag_id="dag", run_id="run", task_id="task", key="key", **extra, **{field: True}) + migrator = get_schema_version_migrator() + + assert field not in migrator.downgrade(message, "2026-06-16").model_dump() + assert migrator.downgrade(message, "2026-10-30").model_dump()[field] is True + + class _IntroduceQueueCapacity(VersionChange): """3026-04-17: introduce queue_capacity.""" @@ -485,6 +506,22 @@ def test_downgrade_strips_arg_bindings_for_previous_version(self, real_migrator, out = real_migrator.downgrade(startup_details, "2026-06-16").model_dump() assert "arg_bindings" not in out["ti_context"] + @pytest.mark.parametrize("version", ["2026-06-16", "2026-10-30"]) + def test_loop_context_follows_supervisor_version(self, real_migrator, startup_details, version): + startup_details.ti_context.loop = LoopContext( + node_id="body", + index=2, + max_iterations=4, + terminal_task_id="body.terminal", + terminal_is_mapped=True, + ) + context = real_migrator.downgrade(startup_details, version).model_dump()["ti_context"] + + if version == "2026-06-16": + assert "loop" not in context + else: + assert context["loop"] == startup_details.ti_context.loop.model_dump() + @pytest.mark.parametrize("version", ["2026-06-16", "2026-10-30"]) def test_task_coordinates_follow_supervisor_version(self, real_migrator, startup_details, version): ti = real_migrator.downgrade(startup_details, version).model_dump()["ti"] diff --git a/task-sdk/tests/task_sdk/execution_time/test_loop.py b/task-sdk/tests/task_sdk/execution_time/test_loop.py new file mode 100644 index 0000000000000..a2485584670bf --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/test_loop.py @@ -0,0 +1,245 @@ +# 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. +from __future__ import annotations + +import pytest + +from airflow.sdk import DAG, BaseOperator, TaskInstanceState, task, task_group +from airflow.sdk.api.datamodels._generated import LoopContext +from airflow.sdk.bases.xcom import BaseXCom +from airflow.sdk.definitions._internal.loop import create_loop +from airflow.sdk.exceptions import AirflowFailException +from airflow.sdk.execution_time import task_runner +from airflow.sdk.execution_time.comms import ( + DeleteXCom, + GetXCom, + GetXComCount, + GetXComSequenceItem, + GetXComSequenceSlice, + SetXCom, + XComCountResponse, + XComResult, + XComSequenceIndexResult, + XComSequenceSliceResult, +) +from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence +from airflow.sdk.execution_time.loop import LoopMaxIterationsExceeded + + +@pytest.fixture +def loop_ti(create_runtime_ti): + def make(*, index=0, max_iterations=3, until=None, map_index=-1): + @task_group + def body(): + BaseOperator(task_id="terminal") + + with DAG("loop_runtime", schedule=None) as dag: + group = create_loop(body, max_iterations=max_iterations, until=until) + ti = create_runtime_ti(task=dag.get_task(group.gate_task_id), map_index=map_index) + ti._ti_context_from_server.loop = LoopContext( + node_id=group.group_id, + index=index, + max_iterations=max_iterations, + terminal_task_id=group.terminal_task_id, + terminal_is_mapped=False, + ) + return ti + + return make + + +def test_first_iteration_previous_does_not_read_xcom(loop_ti, mock_supervisor_comms): + ti = loop_ti(map_index=7) + loop = ti.get_template_context()["loop"] + + assert loop.index == 0 + assert loop.max_iterations == 3 + assert loop.previous is None + assert ti.map_index == 7 + mock_supervisor_comms.send.assert_not_called() + + +@pytest.mark.parametrize("mapped", [False, True]) +def test_body_callable_receives_loop_context_after_unmapping(create_runtime_ti, mapped): + @task + def terminal(value, *, loop, ti): + return value, loop.index, ti.map_index + + @task_group + def body(): + if mapped: + terminal.expand(value=[5]) + else: + terminal(5) + + with DAG("loop_body_runtime", schedule=None) as dag: + group = create_loop(body, max_iterations=4) + operator = dag.get_task(group.terminal_task_id) + if mapped: + operator = operator.unmap({"op_kwargs": {"value": 5}}) + ti = create_runtime_ti(task=operator, map_index=0 if mapped else -1) + ti._ti_context_from_server.loop = LoopContext( + node_id=group.group_id, + index=2, + max_iterations=4, + terminal_task_id=group.terminal_task_id, + terminal_is_mapped=mapped, + ) + + assert operator.execute(ti.get_template_context()) == (5, 2, 0 if mapped else -1) + + +@pytest.mark.parametrize("value", [0, False, [], ""]) +def test_previous_and_current_result_preserve_falsey_values(loop_ti, mock_supervisor_comms, value): + ti = loop_ti(index=1) + mock_supervisor_comms.send.return_value = XComResult(key="return_value", value=value) + loop = ti.get_template_context()["loop"] + + assert loop.previous == value + previous = mock_supervisor_comms.send.call_args.args[0] + assert isinstance(previous, GetXCom) + assert previous.task_id == "body.terminal" + assert previous.previous_iteration is True + assert loop.result == value + assert mock_supervisor_comms.send.call_args.args[0].previous_iteration is False + + +@pytest.mark.parametrize( + ("index", "condition", "decision"), + [(0, None, "continue"), (2, None, "stop"), (0, False, "continue"), (0, True, "stop"), (2, True, "stop")], +) +def test_gate_publishes_successful_decision(loop_ti, mock_supervisor_comms, index, condition, decision): + def until(*, loop): + assert loop.index == index + return condition + + ti = loop_ti(index=index, until=until if condition is not None else None) + context = ti.get_template_context() + + ti.task.execute(context) + + message = mock_supervisor_comms.send.call_args.args[0] + assert isinstance(message, SetXCom) + assert message.key == "_airflow_loop_decision" + assert message.value == decision + + +@pytest.mark.parametrize( + ("raises", "error", "message"), + [(False, LoopMaxIterationsExceeded, "max_iterations"), (True, ValueError, "condition failed")], +) +def test_unsuccessful_gate_does_not_publish_decision(loop_ti, mock_supervisor_comms, raises, error, message): + def until(*, loop): + if raises: + raise ValueError("condition failed") + return False + + ti = loop_ti(index=2, until=until) + with pytest.raises(error, match=message): + ti.task.execute(ti.get_template_context()) + + mock_supervisor_comms.send.assert_not_called() + + +def test_gate_fails_without_retrying_when_the_condition_is_false_at_the_cap(loop_ti, mock_supervisor_comms): + ti = loop_ti(index=2, until=lambda: False) + + with pytest.raises(LoopMaxIterationsExceeded, match="max_iterations=3") as exc_info: + ti.task.execute(ti.get_template_context()) + + assert isinstance(exc_info.value, AirflowFailException) + mock_supervisor_comms.send.assert_not_called() + + +@pytest.mark.parametrize("result", [None, 0, 1, [], "yes"]) +def test_gate_fails_when_until_does_not_return_a_bool(loop_ti, mock_supervisor_comms, result): + ti = loop_ti(index=0, until=lambda: result) + + with pytest.raises(AirflowFailException, match=f"got {type(result).__name__}"): + ti.task.execute(ti.get_template_context()) + + mock_supervisor_comms.send.assert_not_called() + + +@pytest.mark.parametrize("previous_iteration", [False, True]) +@pytest.mark.parametrize("values", [[], [0], [0, False, ""]]) +def test_mapped_terminal_keeps_iteration_across_lazy_reads( + loop_ti, mock_supervisor_comms, previous_iteration, values +): + ti = loop_ti(index=1) + ti._ti_context_from_server.loop.terminal_is_mapped = True + + def respond(message): + assert message.previous_iteration is previous_iteration + if isinstance(message, GetXComCount): + return XComCountResponse(len=len(values)) + if isinstance(message, GetXComSequenceItem): + return XComSequenceIndexResult(root=values[message.offset]) + assert isinstance(message, GetXComSequenceSlice) + return XComSequenceSliceResult(root=values[slice(message.start, message.stop, message.step)]) + + mock_supervisor_comms.send.side_effect = respond + context = ti.get_template_context()["loop"] + result = context.previous if previous_iteration else context.result + + assert isinstance(result, LazyXComSequence) + assert len(result) == len(values) + assert result[:] == values + assert result[::-1] == values[::-1] + if values: + assert result[-1] == values[-1] + + +def test_loop_decision_bypasses_custom_backend(loop_ti, mock_supervisor_comms, mocker): + serialize = mocker.patch("airflow.sdk.execution_time.xcom.XCom.serialize_value", autospec=True) + ti = loop_ti() + + ti.task.execute(ti.get_template_context()) + + serialize.assert_not_called() + assert mock_supervisor_comms.send.call_args.args[0].value == "continue" + + +def test_gate_retry_clears_old_signal_without_custom_backend(loop_ti, mock_supervisor_comms, monkeypatch): + class CustomXCom(BaseXCom): + @classmethod + def purge(cls, xcom, *args): + raise AssertionError("Control metadata must not reach custom backend") + + ti = loop_ti() + ti._ti_context_from_server.xcom_keys_to_clear = ["_airflow_loop_decision"] + monkeypatch.setattr(task_runner, "XCom", CustomXCom) + mock_supervisor_comms.send.side_effect = lambda message: ( + XComResult(key="_airflow_loop_decision", value="stop") if isinstance(message, GetXCom) else None + ) + + state, _, error = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log) + + assert error is None + assert state == TaskInstanceState.SUCCESS + decisions = [ + call.args[0] + for call in mock_supervisor_comms.send.call_args_list + if call.args and isinstance(call.args[0], SetXCom) + ] + assert [decision.value for decision in decisions] == ["continue"] + deletes = [ + call.args[0] + for call in mock_supervisor_comms.send.call_args_list + if call.args and isinstance(call.args[0], DeleteXCom) + ] + assert [delete.key for delete in deletes] == ["_airflow_loop_decision"] diff --git a/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py b/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py index 12a315324621a..9678dd4e09a25 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py +++ b/task-sdk/tests/task_sdk/execution_time/test_request_handlers.py @@ -16,9 +16,12 @@ # under the License. from __future__ import annotations +import json from unittest.mock import MagicMock +import httpx import pytest +from task_sdk import make_client from airflow.sdk.api import client as sdk_client from airflow.sdk.api.datamodels._generated import AssetStateStoreResponse @@ -32,6 +35,10 @@ ErrorResponse, GetAssetStateStoreByName, GetAssetStateStoreByUri, + GetXCom, + GetXComCount, + GetXComSequenceItem, + GetXComSequenceSlice, SetAssetStateStoreByName, SetAssetStateStoreByUri, ) @@ -42,6 +49,10 @@ handle_delete_asset_state_store_by_uri, handle_get_asset_state_store_by_name, handle_get_asset_state_store_by_uri, + handle_get_xcom, + handle_get_xcom_count, + handle_get_xcom_sequence_item, + handle_get_xcom_sequence_slice, handle_set_asset_state_store_by_name, handle_set_asset_state_store_by_uri, ) @@ -59,6 +70,39 @@ def client_ssl_cache(): sdk_client.Client._get_ssl_context_cached.cache_clear() +@pytest.mark.parametrize( + ("message", "handler", "extra", "response", "field"), + [ + (GetXCom, handle_get_xcom, {}, {"key": "key", "value": 0}, "previous_iteration"), + (GetXComCount, handle_get_xcom_count, {}, None, "previous_iteration"), + (GetXComSequenceItem, handle_get_xcom_sequence_item, {"offset": -1}, 0, "previous_iteration"), + ( + GetXComSequenceSlice, + handle_get_xcom_sequence_slice, + {"start": None, "stop": None, "step": -1}, + [], + "previous_iteration", + ), + ], +) +def test_loop_selectors_survive_comms_and_http_transport( + message, handler, extra, response, field, client_ssl_cache +): + msg = message(dag_id="dag", run_id="run", task_id="task", key="key", **extra, **{field: True}) + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(200, content=json.dumps(response), headers={"Content-Range": "map_indexes 1"}) + + handler( + make_client(transport=httpx.MockTransport(respond)), + message.model_validate_json(msg.model_dump_json()), + ) + + assert requests[0].url.params[field] == "true" + + def test_get_asset_state_store_by_name_wraps_response_as_result(client): client.asset_state_store.get.return_value = AssetStateStoreResponse(value="2026-01-01") diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index d109dd4ffd797..0f507a87f8e47 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -2441,6 +2441,7 @@ class RequestTestCase: client_mock=ClientMock( method_path="xcoms.get", args=("test_dag", "test_run", "test_task", "test_key", None, False), + kwargs={"previous_iteration": False}, response=XComResult(key="test_key", value="test_value"), ), expected_body={"key": "test_key", "value": "test_value", "type": "XComResult"}, @@ -2453,6 +2454,7 @@ class RequestTestCase: client_mock=ClientMock( method_path="xcoms.get", args=("test_dag", "test_run", "test_task", "test_key", 2, False), + kwargs={"previous_iteration": False}, response=XComResult(key="test_key", value="test_value"), ), expected_body={"key": "test_key", "value": "test_value", "type": "XComResult"}, @@ -2463,6 +2465,7 @@ class RequestTestCase: client_mock=ClientMock( method_path="xcoms.get", args=("test_dag", "test_run", "test_task", "test_key", None, False), + kwargs={"previous_iteration": False}, response=XComResult(key="test_key", value=None, type="XComResult"), ), expected_body={"key": "test_key", "value": None, "type": "XComResult"}, @@ -2479,6 +2482,7 @@ class RequestTestCase: client_mock=ClientMock( method_path="xcoms.get", args=("test_dag", "test_run", "test_task", "test_key", None, True), + kwargs={"previous_iteration": False}, response=XComResult(key="test_key", value=None, type="XComResult"), ), expected_body={"key": "test_key", "value": None, "type": "XComResult"}, @@ -2501,7 +2505,10 @@ class RequestTestCase: '{"key": "test_key", "value": {"key2": "value2"}}', None, ), - kwargs={"dag_result": False, "mapped_length": None}, + kwargs={ + "dag_result": False, + "mapped_length": None, + }, response=OKResponse(ok=True), ), test_id="set_xcom", @@ -2525,7 +2532,10 @@ class RequestTestCase: '{"key": "test_key", "value": {"key2": "value2"}}', 2, ), - kwargs={"dag_result": False, "mapped_length": None}, + kwargs={ + "dag_result": False, + "mapped_length": None, + }, response=OKResponse(ok=True), ), test_id="set_xcom_with_map_index", @@ -2550,7 +2560,10 @@ class RequestTestCase: '{"key": "test_key", "value": {"key2": "value2"}}', 2, ), - kwargs={"dag_result": False, "mapped_length": 3}, + kwargs={ + "dag_result": False, + "mapped_length": 3, + }, response=OKResponse(ok=True), ), test_id="set_xcom_with_map_index_and_mapped_length", @@ -2574,7 +2587,10 @@ class RequestTestCase: '{"key": "test_key", "value": {"key2": "value2"}}', None, ), - kwargs={"dag_result": True, "mapped_length": None}, + kwargs={ + "dag_result": True, + "mapped_length": None, + }, response=OKResponse(ok=True), ), test_id="set_xcom_with_dag_result", @@ -3570,6 +3586,7 @@ class RequestTestCase: client_mock=ClientMock( method_path="xcoms.get_sequence_item", args=("test_dag", "test_run", "test_task", "test_key", 0), + kwargs={"previous_iteration": False}, response=XComSequenceIndexResult(root="test_value"), ), test_id="get_xcom_seq_item", @@ -3586,6 +3603,7 @@ class RequestTestCase: client_mock=ClientMock( method_path="xcoms.get_sequence_item", args=("test_dag", "test_run", "test_task", "test_key", 2), + kwargs={"previous_iteration": False}, response=ErrorResponse(error=ErrorType.XCOM_NOT_FOUND), ), test_id="get_xcom_seq_item_not_found", @@ -3605,6 +3623,7 @@ class RequestTestCase: client_mock=ClientMock( method_path="xcoms.get_sequence_slice", args=("test_dag", "test_run", "test_task", "test_key", None, None, None, False), + kwargs={"previous_iteration": False}, response=XComSequenceSliceResult(root=["foo", "bar"]), ), test_id="get_xcom_seq_slice", @@ -3667,6 +3686,7 @@ class RequestTestCase: client_mock=ClientMock( method_path="xcoms.head", args=("test_dag", "test_run", "test_task", "test_key"), + kwargs={"previous_iteration": False}, response=XComCountResponse(len=5), ), test_id="get_xcom_count", diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 66fa2411082d4..8d4ade1a3401b 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -64,6 +64,7 @@ dag as dag_decorator, get_current_context, task as task_decorator, + task_group, timezone, ) from airflow.sdk._shared.observability.metrics.base_stats_logger import StatsLogger @@ -83,6 +84,7 @@ from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter from airflow.sdk.coordinators._subprocess import SubprocessCoordinator +from airflow.sdk.definitions._internal.loop import create_loop from airflow.sdk.definitions._internal.types import NOTSET, SET_DURING_EXECUTION, is_arg_set from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetUniqueKey, AssetUriRef, Dataset, Model from airflow.sdk.definitions.param import DagParam @@ -177,6 +179,7 @@ VariableAccessor, _wrap_external_ref, ) +from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence from airflow.sdk.execution_time.task_runner import ( RuntimeTaskInstance, TaskRunnerMarker, @@ -3246,6 +3249,67 @@ def mock_send_side_effect(*args, **kwargs): ), ) + @staticmethod + def _make_loop_dag(): + @task_group + def body(): + BaseOperator(task_id="first") >> BaseOperator(task_id="second") + + with DAG("loop_dag", schedule=None) as dag: + create_loop(body, max_iterations=3) >> BaseOperator(task_id="after") + return dag + + def test_xcom_pull_of_a_loop_task_from_outside_the_loop_returns_every_iteration( + self, create_runtime_ti, mock_supervisor_comms + ): + dag = self._make_loop_dag() + runtime_ti = create_runtime_ti(task=dag.get_task("after")) + mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=[1, 2, 3]) + + result = runtime_ti.xcom_pull(key="test_key", task_ids="body.first") + + assert isinstance(result, LazyXComSequence) + mock_supervisor_comms.send.assert_not_called() + assert result[:] == [1, 2, 3] + sent = mock_supervisor_comms.send.call_args.args[0] + assert (sent.task_id, sent.key, sent.previous_iteration) == ("body.first", "test_key", False) + + @pytest.mark.parametrize( + ("task_id", "pull_kwargs"), + [ + pytest.param("body.second", {"task_ids": "body.first"}, id="caller inside the loop"), + pytest.param("after", {"task_ids": "body.first", "include_prior_dates": True}, id="prior dates"), + pytest.param("after", {"task_ids": "body.first", "run_id": "other_run"}, id="other run"), + pytest.param("after", {"task_ids": "body.first", "map_indexes": -1}, id="explicit map index"), + pytest.param("after", {"task_ids": ["body.first"]}, id="several task ids"), + pytest.param("after", {"task_ids": "removed"}, id="task missing from the dag"), + pytest.param("after", {"task_ids": "after"}, id="task outside any loop"), + ], + ) + def test_xcom_pull_keeps_pulling_values_unless_an_outside_caller_reads_a_loop_task( + self, create_runtime_ti, mock_supervisor_comms, task_id, pull_kwargs + ): + dag = self._make_loop_dag() + runtime_ti = create_runtime_ti(task=dag.get_task(task_id)) + mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=[1, 2, 3]) + + with patch.object(XCom, "get_one", return_value=7): + result = runtime_ti.xcom_pull(key="test_key", **pull_kwargs) + + assert not isinstance(result, LazyXComSequence) + + @pytest.mark.asyncio + async def test_axcom_pull_of_a_loop_task_from_outside_the_loop_returns_a_list_for_one_iteration( + self, create_runtime_ti, mock_supervisor_comms + ): + dag = self._make_loop_dag() + runtime_ti = create_runtime_ti(task=dag.get_task("after")) + mock_supervisor_comms.asend.return_value = XComSequenceSliceResult(root=["only"]) + + result = await runtime_ti.axcom_pull(key="test_key", task_ids="body.first") + + assert result == ["only"] + def test_get_param_from_context( self, mocked_parse, make_ti_context, mock_supervisor_comms, create_runtime_ti ): diff --git a/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json index b0d5797b788c8..7e09a736de793 100644 --- a/ts-sdk/schema/supervisor-schema.json +++ b/ts-sdk/schema/supervisor-schema.json @@ -2836,6 +2836,11 @@ }, "GetXCom": { "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -2888,6 +2893,11 @@ "GetXComCount": { "description": "Get the number of (mapped) XCom values available.", "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -2922,6 +2932,11 @@ }, "GetXComSequenceItem": { "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -2961,6 +2976,11 @@ }, "GetXComSequenceSlice": { "properties": { + "previous_iteration": { + "default": false, + "title": "Previous Iteration", + "type": "boolean" + }, "key": { "title": "Key", "type": "string" @@ -3222,6 +3242,42 @@ "title": "LazyDeserializedDAG", "type": "object" }, + "LoopContext": { + "description": "Pinned loop definition and enclosing iteration for a task execution.", + "properties": { + "node_id": { + "title": "Node Id", + "type": "string" + }, + "index": { + "minimum": 0, + "title": "Index", + "type": "integer" + }, + "max_iterations": { + "exclusiveMinimum": 0, + "title": "Max Iterations", + "type": "integer" + }, + "terminal_task_id": { + "title": "Terminal Task Id", + "type": "string" + }, + "terminal_is_mapped": { + "title": "Terminal Is Mapped", + "type": "boolean" + } + }, + "required": [ + "node_id", + "index", + "max_iterations", + "terminal_task_id", + "terminal_is_mapped" + ], + "title": "LoopContext", + "type": "object" + }, "MaskSecret": { "description": "Add a new value to be redacted in task logs.", "properties": { @@ -5154,6 +5210,17 @@ "dag_run": { "$ref": "#/$defs/DagRun" }, + "loop": { + "anyOf": [ + { + "$ref": "#/$defs/LoopContext" + }, + { + "type": "null" + } + ], + "default": null + }, "task_reschedule_count": { "default": 0, "title": "Task Reschedule Count", diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 1d785b4d47600..31196f868a5cb 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -231,6 +231,11 @@ export type TaskInstanceState = | "skipped" | "deferred" | "awaiting_input"; +export type NodeId = string; +export type Index = number; +export type MaxIterations = number; +export type TerminalTaskId = string; +export type TerminalIsMapped = boolean; export type TaskRescheduleCount = number; export type MaxTries = number; export type Key = string; @@ -473,6 +478,7 @@ export type Prefix = string | null; export type Limit2 = number; export type Offset = number; export type Type49 = "GetVariableKeys"; +export type PreviousIteration = boolean; export type Key10 = string; export type DagId16 = string; export type RunId9 = string; @@ -480,17 +486,20 @@ export type TaskId4 = string; export type MapIndex5 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; +export type PreviousIteration1 = boolean; export type Key11 = string; export type DagId17 = string; export type RunId10 = string; export type TaskId5 = string; export type Type51 = "GetXComCount"; +export type PreviousIteration2 = boolean; export type Key12 = string; export type DagId18 = string; export type RunId11 = string; export type TaskId6 = string; export type Offset1 = number; export type Type52 = "GetXComSequenceItem"; +export type PreviousIteration3 = boolean; export type Key13 = string; export type DagId19 = string; export type RunId12 = string; @@ -1038,6 +1047,7 @@ export interface TaskCallbackRequest { */ export interface TIRunContext { dag_run: DagRun; + loop?: LoopContext | null; task_reschedule_count?: TaskRescheduleCount; max_tries: MaxTries; variables?: Variables; @@ -1050,6 +1060,19 @@ export interface TIRunContext { arg_bindings?: ArgBindings; multi_team?: MultiTeam; } +/** + * Pinned loop definition and enclosing iteration for a task execution. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "LoopContext". + */ +export interface LoopContext { + node_id: NodeId; + index: Index; + max_iterations: MaxIterations; + terminal_task_id: TerminalTaskId; + terminal_is_mapped: TerminalIsMapped; +} /** * Variable schema for responses with fields that are needed for Runtime. * @@ -1514,6 +1537,7 @@ export interface GetVariableKeys { * via the `definition` "GetXCom". */ export interface GetXCom { + previous_iteration?: PreviousIteration; key: Key10; dag_id: DagId16; run_id: RunId9; @@ -1529,6 +1553,7 @@ export interface GetXCom { * via the `definition` "GetXComCount". */ export interface GetXComCount { + previous_iteration?: PreviousIteration1; key: Key11; dag_id: DagId17; run_id: RunId10; @@ -1540,6 +1565,7 @@ export interface GetXComCount { * via the `definition` "GetXComSequenceItem". */ export interface GetXComSequenceItem { + previous_iteration?: PreviousIteration2; key: Key12; dag_id: DagId18; run_id: RunId11; @@ -1552,6 +1578,7 @@ export interface GetXComSequenceItem { * via the `definition` "GetXComSequenceSlice". */ export interface GetXComSequenceSlice { + previous_iteration?: PreviousIteration3; key: Key13; dag_id: DagId19; run_id: RunId12;