You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2026/10/06 15:02:00 UTC
(airflow) branch task-loops-stack-8 created (now 09978c6ac6f)
This is an automated email from the ASF dual-hosted git repository.
ashb pushed a change to branch task-loops-stack-8
in repository https://gitbox.apache.org/repos/asf/airflow.git
at 09978c6ac6f Clear selected loop work and regenerate only what was replaced
This branch includes the following new commits:
new 09978c6ac6f Clear selected loop work and regenerate only what was replaced
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(airflow) 01/01: Clear selected loop work and regenerate only what was replaced
Posted by as...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch task-loops-stack-8
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 09978c6ac6fe3a32fc6ef662d14a9f051094faa8
Author: Ash Berlin-Taylor <as...@apache.org>
AuthorDate: Sun Oct 4 08:57:42 2026 +0100
Clear selected loop work and regenerate only what was replaced
Clearing one task in a loop must not disturb unrelated work in the same
iteration, and clearing a gate must not let a rerun gate redo a progression that
later iterations already followed. The clear API therefore takes two controls
that default on: clear downstream, and clear later loop iterations. Rerunning a
gate while later iterations remain leaves them alone, because a second decision
for the same iteration would fork history with nothing to attribute it to.
Replaced generated work leaves the live set by being retired in place, so its
tries, logs, XCom and notes stay attached to the attempt that produced them,
and the work that replaces it is created in a new region instead of reoccupying
the old coordinates. Ordinary tasks still get a successor try in their existing
region. A pre-region mapped expansion has no region record, so a whole-task
clear mints one rather than rewriting the sentinel rows.
We cannot retire a row that is still terminating from the clear transaction
alone: the supervisor's completion report decides how the row ends. Clears,
executor events, the heartbeat scan and gate completion therefore take the
DagRun lock before any task instance lock, so they cannot interleave in opposite
orders and deadlock. Rows that never started are retired without being stamped
FAILED, so their state still says they never ran.
When the scheduler skips an executor event's task instance row because another
transaction holds its lock, it requeues the event, as it does when the DagRun is
locked. Dropping it would leave the attempt unreported. The event is discarded
only once the row is gone or retired, and a requeued event is not counted as
processed, so a held lock neither inflates the metrics nor keeps the scheduler
loop from sleeping.
---
.../core_api/datamodels/task_instances.py | 30 +
.../api_fastapi/core_api/openapi/_private_ui.yaml | 9 +
.../core_api/openapi/v2-rest-api-generated.yaml | 157 ++++
.../core_api/routes/public/task_instances.py | 294 +++++--
.../core_api/services/public/dag_run.py | 12 +-
.../core_api/services/public/task_coordinates.py | 113 +++
.../core_api/services/public/task_instances.py | 411 ++++++++-
.../execution_api/routes/task_instances.py | 19 +-
.../src/airflow/cli/commands/dag_command.py | 4 +-
.../src/airflow/jobs/scheduler_job_runner.py | 70 +-
airflow-core/src/airflow/models/dagbag.py | 21 -
airflow-core/src/airflow/models/dagrun.py | 156 +++-
airflow-core/src/airflow/models/loop_clear.py | 449 ++++++++++
airflow-core/src/airflow/models/taskinstance.py | 115 ++-
.../serialization/definitions/baseoperator.py | 2 +
.../src/airflow/serialization/definitions/dag.py | 20 +-
.../src/airflow/ti_deps/deps/loop_archival_dep.py | 32 +
.../src/airflow/ui/openapi-gen/queries/queries.ts | 48 +-
.../airflow/ui/openapi-gen/requests/schemas.gen.ts | 89 +-
.../ui/openapi-gen/requests/services.gen.ts | 36 +-
.../airflow/ui/openapi-gen/requests/types.gen.ts | 30 +
.../TaskInstance/MarkTaskInstanceAsDialog.test.tsx | 2 +
.../ui/src/components/TaskInstanceTooltip.test.tsx | 2 +
.../core_api/routes/public/test_dag_run.py | 34 +
.../core_api/routes/public/test_task_instances.py | 967 ++++++++++++++++++++-
.../versions/head/test_task_instances.py | 295 ++++++-
.../tests/unit/cli/commands/test_dag_command.py | 11 +-
.../unit/cli/commands/test_partition_command.py | 8 +-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 139 ++-
airflow-core/tests/unit/models/test_cleartasks.py | 84 +-
airflow-core/tests/unit/models/test_dag.py | 2 +-
airflow-core/tests/unit/models/test_dagrun.py | 4 +-
.../tests/unit/models/test_dynamic_region.py | 12 +-
airflow-core/tests/unit/models/test_loop_clear.py | 947 ++++++++++++++++++++
.../tests/unit/models/test_taskinstance.py | 97 +++
.../src/airflowctl/api/datamodels/generated.py | 38 +
.../tests/airflow_ctl/api/test_operations.py | 4 +
.../airflow_ctl/ctl/commands/test_task_command.py | 6 +
.../tests/unit/openlineage/utils/test_utils.py | 10 +-
39 files changed, 4574 insertions(+), 205 deletions(-)
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
index 1e7564afe48..4855479dc15 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
@@ -57,6 +57,8 @@ class TaskInstanceResponse(BaseModel):
dag_id: str
run_id: str = Field(alias="dag_run_id")
map_index: int
+ region_id: UUID
+ region_index: int
logical_date: datetime | None
run_after: datetime
start_date: datetime | None
@@ -239,6 +241,19 @@ class ClearTaskInstancesBody(StrictBaseModel):
"The group's tasks are resolved on the server from the dag structure, so all of them are "
"targeted regardless of how many there are.",
)
+ task_instance_ids: list[UUID] | None = Field(
+ default=None,
+ min_length=1,
+ description="Exact current executions in one DAG run. Mutually exclusive with task_ids and task_group_id.",
+ )
+ whole_expansion_ids: list[UUID] = Field(
+ default_factory=list,
+ description="Selected execution IDs whose entire mapped expansion should be cleared.",
+ )
+ include_later_loop_iterations: bool = Field(
+ default=True,
+ description="Clear later iterations when the selection includes a loop gate.",
+ )
dag_run_id: str | None = None
include_upstream: bool = False
include_downstream: bool = False
@@ -281,11 +296,26 @@ class ClearTaskInstancesBody(StrictBaseModel):
raise ValueError("Only one of task_ids or task_group_id may be provided")
return data
+ @model_validator(mode="after")
+ def validate_execution_selection(self):
+ if self.task_instance_ids is not None:
+ if self.task_ids is not None or self.task_group_id is not None:
+ raise ValueError("task_instance_ids cannot be combined with task_ids or task_group_id")
+ if self.dag_run_id is None or self.include_past or self.include_future:
+ raise ValueError(
+ "task_instance_ids requires one dag_run_id without include_past or include_future"
+ )
+ if not set(self.whole_expansion_ids) <= set(self.task_instance_ids or ()):
+ raise ValueError("whole_expansion_ids must be a subset of task_instance_ids")
+ return self
+
class PatchTaskInstanceBody(StrictBaseModel):
"""Request body for patching task instance state."""
new_state: TaskInstanceState | None = None
+ region_id: UUID | None = None
+ region_index: int | None = None
note: Annotated[str, StringConstraints(max_length=1000)] | None = None
include_upstream: bool = False
include_downstream: bool = False
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
index de12f997600..8a7236631fa 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
+++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
@@ -4842,6 +4842,13 @@ components:
map_index:
type: integer
title: Map Index
+ region_id:
+ type: string
+ format: uuid
+ title: Region Id
+ region_index:
+ type: integer
+ title: Region Index
logical_date:
anyOf:
- type: string
@@ -4993,6 +5000,8 @@ components:
- dag_id
- dag_run_id
- map_index
+ - region_id
+ - region_index
- logical_date
- run_after
- start_date
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
index 757e0e06ff6..625d3f28760 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
+++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
@@ -7534,6 +7534,23 @@ paths:
type: string
- type: 'null'
title: Update Mask
+ - name: region_id
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: string
+ format: uuid
+ - type: 'null'
+ title: Region Id
+ - name: region_index
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: integer
+ - type: 'null'
+ title: Region Index
requestBody:
required: true
content:
@@ -8499,6 +8516,23 @@ paths:
type: string
- type: 'null'
title: Update Mask
+ - name: region_id
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: string
+ format: uuid
+ - type: 'null'
+ title: Region Id
+ - name: region_index
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: integer
+ - type: 'null'
+ title: Region Index
requestBody:
required: true
content:
@@ -9537,6 +9571,23 @@ paths:
type: string
- type: 'null'
title: Update Mask
+ - name: region_id
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: string
+ format: uuid
+ - type: 'null'
+ title: Region Id
+ - name: region_index
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: integer
+ - type: 'null'
+ title: Region Index
requestBody:
required: true
content:
@@ -9615,6 +9666,23 @@ paths:
schema:
type: string
title: Group Id
+ - name: region_id
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: string
+ format: uuid
+ - type: 'null'
+ title: Region Id
+ - name: region_index
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: integer
+ - type: 'null'
+ title: Region Index
requestBody:
required: true
content:
@@ -9705,6 +9773,23 @@ paths:
type: string
- type: 'null'
title: Update Mask
+ - name: region_id
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: string
+ format: uuid
+ - type: 'null'
+ title: Region Id
+ - name: region_index
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: integer
+ - type: 'null'
+ title: Region Index
requestBody:
required: true
content:
@@ -9795,6 +9880,23 @@ paths:
type: string
- type: 'null'
title: Update Mask
+ - name: region_id
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: string
+ format: uuid
+ - type: 'null'
+ title: Region Id
+ - name: region_index
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: integer
+ - type: 'null'
+ title: Region Index
requestBody:
required: true
content:
@@ -12802,6 +12904,17 @@ components:
anyOf:
- $ref: '#/components/schemas/TaskInstanceState'
- type: 'null'
+ region_id:
+ anyOf:
+ - type: string
+ format: uuid
+ - type: 'null'
+ title: Region Id
+ region_index:
+ anyOf:
+ - type: integer
+ - type: 'null'
+ title: Region Index
note:
anyOf:
- type: string
@@ -13181,6 +13294,30 @@ components:
description: Clear every task in this task group. Mutually exclusive with
`task_ids`. The group's tasks are resolved on the server from the dag
structure, so all of them are targeted regardless of how many there are.
+ task_instance_ids:
+ anyOf:
+ - items:
+ type: string
+ format: uuid
+ type: array
+ minItems: 1
+ - type: 'null'
+ title: Task Instance Ids
+ description: Exact current executions in one DAG run. Mutually exclusive
+ with task_ids and task_group_id.
+ whole_expansion_ids:
+ items:
+ type: string
+ format: uuid
+ type: array
+ title: Whole Expansion Ids
+ description: Selected execution IDs whose entire mapped expansion should
+ be cleared.
+ include_later_loop_iterations:
+ type: boolean
+ title: Include Later Loop Iterations
+ description: Clear later iterations when the selection includes a loop gate.
+ default: true
dag_run_id:
anyOf:
- type: string
@@ -16086,6 +16223,17 @@ components:
anyOf:
- $ref: '#/components/schemas/TaskInstanceState'
- type: 'null'
+ region_id:
+ anyOf:
+ - type: string
+ format: uuid
+ - type: 'null'
+ title: Region Id
+ region_index:
+ anyOf:
+ - type: integer
+ - type: 'null'
+ title: Region Index
note:
anyOf:
- type: string
@@ -16950,6 +17098,13 @@ components:
map_index:
type: integer
title: Map Index
+ region_id:
+ type: string
+ format: uuid
+ title: Region Id
+ region_index:
+ type: integer
+ title: Region Index
logical_date:
anyOf:
- type: string
@@ -17101,6 +17256,8 @@ components:
- dag_id
- dag_run_id
- map_index
+ - region_id
+ - region_index
- logical_date
- run_after
- start_date
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
index 1f40a0c73e7..6bddcf3d882 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
@@ -17,12 +17,14 @@
from __future__ import annotations
+from collections import defaultdict
from collections.abc import Sequence
from typing import Annotated, Literal, cast
+from uuid import UUID
import structlog
from fastapi import Depends, HTTPException, Query, status
-from sqlalchemy import select
+from sqlalchemy import select, tuple_
from sqlalchemy.orm import joinedload
from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
@@ -34,7 +36,6 @@ from airflow.api_fastapi.common.cursors import (
)
from airflow.api_fastapi.common.dagbag import (
DagBagDep,
- get_dag_for_run,
get_dag_for_run_or_latest_version,
get_latest_version_of_dag,
resolve_run_on_latest_version,
@@ -104,27 +105,30 @@ from airflow.api_fastapi.core_api.datamodels.task_instances import (
)
from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
from airflow.api_fastapi.core_api.security import GetUserDep, ReadableTIFilterDep, requires_access_dag
+from airflow.api_fastapi.core_api.services.public.task_coordinates import task_coordinate_response
from airflow.api_fastapi.core_api.services.public.task_instances import (
BulkTaskInstanceService,
_discard_task_state_store,
_get_task_group_task_ids,
- _get_task_group_task_instances,
_patch_task_group_state,
_patch_task_instance_note,
_patch_task_instance_state,
_patch_ti_group_validate_request,
_patch_ti_validate_request,
_reload_tis_with_rendered_fields,
+ patch_region_selection,
)
from airflow.api_fastapi.logging.decorators import action_logging
from airflow.exceptions import AirflowClearRunningTaskException, TaskNotFound
from airflow.models import DagRun
+from airflow.models.loop_clear import LoopClearScope, apply_loop_clear_scope, select_loop_clear_scope
from airflow.models.renderedtifields import load_legacy_rendered_fields
+from airflow.models.task_coordinates import TaskCoordinateResolver, enclosing_loop
from airflow.models.taskinstance import TaskInstance as TI, clear_task_instances
from airflow.ti_deps.dep_context import DepContext
from airflow.ti_deps.dependencies_deps import SCHEDULER_QUEUED_DEPS
from airflow.utils.db import get_query_count
-from airflow.utils.state import DagRunState, TaskInstanceState
+from airflow.utils.state import DagRunState, State, TaskInstanceState
log = structlog.get_logger(__name__)
@@ -132,6 +136,13 @@ task_instances_router = AirflowRouter(tags=["Task Instance"], prefix="/dags/{dag
task_instances_prefix = "/dagRuns/{dag_run_id}/taskInstances"
+def _coordinate_resolver(dag_bag: DagBagDep, session: SessionDep) -> TaskCoordinateResolver:
+ return TaskCoordinateResolver(dag_bag, session)
+
+
+CoordinateResolverDep = Annotated[TaskCoordinateResolver, Depends(_coordinate_resolver)]
+
+
@task_instances_router.get(
task_instances_prefix + "/{task_id}",
responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
@@ -873,6 +884,7 @@ def post_clear_task_instances(
body: ClearTaskInstancesBody,
session: SessionDep,
user: GetUserDep,
+ resolver: CoordinateResolverDep,
) -> TaskInstanceCollectionResponse:
"""Clear task instances."""
dag = get_latest_version_of_dag(dag_bag, dag_id, session)
@@ -896,7 +908,7 @@ def post_clear_task_instances(
error_message = f"Dag Run id {dag_run_id} not found in dag {dag_id}"
raise HTTPException(status.HTTP_404_NOT_FOUND, error_message)
# Get the specific dag version:
- dag = get_dag_for_run(dag_bag, dag_run, session)
+ dag = get_dag_for_run_or_latest_version(dag_bag, dag_run, dag_id, session)
if (past or future) and dag_run.logical_date is None:
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
@@ -919,7 +931,21 @@ def post_clear_task_instances(
_get_task_group_task_ids(dag_id, body.task_group_id, dag),
)
- if (task_markers_to_clear := body.task_ids) is not None:
+ gate_query = select(TI.id).join(TI.dag_run).where(TI.dag_id == dag_id, TI.operator == "LoopGateOperator")
+ if dag_run_id is not None and not (past or future):
+ gate_query = gate_query.where(TI.run_id == dag_run_id)
+ else:
+ if body.start_date is not None:
+ gate_query = gate_query.where(DagRun.logical_date >= body.start_date)
+ if body.end_date is not None:
+ gate_query = gate_query.where(DagRun.logical_date <= body.end_date)
+ loop_aware = (
+ body.task_instance_ids is not None
+ or any(enclosing_loop(task) for task in dag.tasks)
+ or session.scalar(gate_query.limit(1)) is not None
+ )
+ task_markers_to_clear = body.task_ids
+ if task_markers_to_clear is not None and not loop_aware:
mapped_tasks_tuples = {t for t in task_markers_to_clear if isinstance(t, tuple)}
# Unmapped tasks are expressed in their task_ids (without map_indexes)
normal_task_ids = {t for t in task_markers_to_clear if not isinstance(t, tuple)}
@@ -963,42 +989,162 @@ def post_clear_task_instances(
*((t, m) for t, m in mapped_tasks_tuples if t not in normal_task_ids),
]
- task_instances: Sequence[TI]
- if dag_run_id is not None and not (past or future):
- # Use run_id-based clearing when we have a specific dag_run_id and not using past/future
- task_instances = dag.clear(
- dry_run=True,
- task_ids=task_markers_to_clear,
- run_id=dag_run_id,
- session=session,
- run_on_latest_version=resolved_run_on_latest,
- only_failed=body.only_failed,
- only_running=body.only_running,
- )
- else:
- # Use date-based clearing when no dag_run_id or when past/future is specified
- task_instances = dag.clear(
+ clear_markers = task_markers_to_clear
+ if loop_aware and task_markers_to_clear is not None:
+ clear_markers = [marker if isinstance(marker, str) else marker[0] for marker in task_markers_to_clear]
+
+ def select_candidates() -> Sequence[TI]:
+ if body.task_instance_ids is not None:
+ return session.scalars(
+ select(TI).where(
+ TI.dag_id == dag_id,
+ TI.run_id == dag_run_id,
+ TI.working_set.is_(True),
+ TI.id.in_(body.task_instance_ids),
+ )
+ ).all()
+ if dag_run_id is not None and not (past or future):
+ return dag.clear(
+ dry_run=True,
+ task_ids=clear_markers,
+ run_id=dag_run_id,
+ session=session,
+ run_on_latest_version=resolved_run_on_latest,
+ only_failed=body.only_failed and not loop_aware,
+ only_running=body.only_running and not loop_aware,
+ )
+ return dag.clear(
dry_run=True,
- task_ids=task_markers_to_clear,
+ task_ids=clear_markers,
start_date=body.start_date,
end_date=body.end_date,
session=session,
run_on_latest_version=resolved_run_on_latest,
- only_failed=body.only_failed,
- only_running=body.only_running,
+ only_failed=body.only_failed and not loop_aware,
+ only_running=body.only_running and not loop_aware,
)
- if not dry_run:
+ task_instances = select_candidates()
+ if body.task_instance_ids is not None and {ti.id for ti in task_instances} != set(body.task_instance_ids):
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND, "Selected task execution is not current in this DAG run"
+ )
+
+ if task_instances:
+ if not dry_run:
+ run_keys = {(ti.dag_id, ti.run_id) for ti in task_instances}
+ session.scalars(
+ select(DagRun)
+ .where(tuple_(DagRun.dag_id, DagRun.run_id).in_(run_keys))
+ .order_by(DagRun.dag_id, DagRun.run_id)
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ ).all()
+ if body.task_instance_ids is None:
+ task_instances = [ti for ti in select_candidates() if (ti.dag_id, ti.run_id) in run_keys]
+ task_instances = session.scalars(
+ select(TI)
+ .where(TI.id.in_([ti.id for ti in task_instances]))
+ .execution_options(populate_existing=True)
+ ).all()
+
+ scopes: list[LoopClearScope] = []
+ if loop_aware and task_markers_to_clear is not None:
+ task_instances = [
+ ti
+ for ti in task_instances
+ if ti.task_id in task_markers_to_clear
+ or (ti.task_id, resolver.public_map_index(ti)) in task_markers_to_clear
+ ]
+ if body.task_instance_ids is not None:
+ if {ti.id for ti in task_instances} != set(body.task_instance_ids):
+ raise HTTPException(status.HTTP_409_CONFLICT, "Selected task execution changed before clearing")
+ if loop_aware:
+ selections: dict[tuple[str, str], list[TI]] = defaultdict(list)
+ for ti in task_instances:
+ selections[(ti.dag_id, ti.run_id)].append(ti)
try:
- task_instances = clear_task_instances(
- task_instances,
- session,
- DagRunState.QUEUED if reset_dag_runs else False,
- run_on_latest_version=resolved_run_on_latest,
- prevent_running_task=body.prevent_running_task,
+ for selected in selections.values():
+ selected_ids = {ti.id for ti in selected}
+ scope = select_loop_clear_scope(
+ selected,
+ whole_expansion_ids=set(body.whole_expansion_ids) & selected_ids,
+ upstream=body.include_upstream,
+ downstream=body.include_downstream,
+ later_loop_iterations=body.include_later_loop_iterations
+ and not (body.only_failed or body.only_running),
+ session=session,
+ )
+ if body.only_failed or body.only_running:
+ eligible = session.scalars(
+ select(TI).where(
+ TI.id.in_(scope.retry_ids),
+ TI.state.in_(State.failed_states)
+ if body.only_failed
+ else TI.state == State.RUNNING,
+ )
+ ).all()
+ scope = select_loop_clear_scope(
+ eligible,
+ downstream=False,
+ later_loop_iterations=body.include_later_loop_iterations,
+ session=session,
+ )
+ scopes.append(scope)
+ except ValueError as error:
+ raise HTTPException(status.HTTP_409_CONFLICT, str(error)) from error
+ task_instances = session.scalars(
+ select(TI)
+ .where(
+ TI.id.in_({identity for scope in scopes for identity in scope.retry_ids | scope.archive_ids})
+ )
+ .execution_options(populate_existing=True)
+ ).all()
+
+ if not dry_run:
+ whole_task_keys = {
+ (ti.dag_id, ti.run_id, ti.task_id)
+ for ti in task_instances
+ if not body.only_failed
+ and not body.only_running
+ and body.task_instance_ids is None
+ and (task_markers_to_clear is None or ti.task_id in task_markers_to_clear)
+ }
+ if not body.only_failed and not body.only_running:
+ whole_task_keys.update(
+ (ti.dag_id, ti.run_id, ti.task_id)
+ for ti in task_instances
+ if ti.id in body.whole_expansion_ids
)
+ try:
+ if loop_aware:
+ cleared: list[TI] = []
+ for scope in scopes:
+ cleared += apply_loop_clear_scope(
+ scope,
+ session=session,
+ later_loop_iterations=body.include_later_loop_iterations,
+ dag_run_state=DagRunState.QUEUED if reset_dag_runs else False,
+ run_on_latest_version=resolved_run_on_latest,
+ prevent_running_task=body.prevent_running_task,
+ whole_task_keys=whole_task_keys,
+ )
+ task_instances = cleared
+ else:
+ task_instances = clear_task_instances(
+ list(task_instances),
+ session,
+ DagRunState.QUEUED if reset_dag_runs else False,
+ run_on_latest_version=resolved_run_on_latest,
+ prevent_running_task=body.prevent_running_task,
+ whole_task_keys=whole_task_keys,
+ )
except AirflowClearRunningTaskException as e:
raise HTTPException(status.HTTP_409_CONFLICT, str(e)) from e
+ except ValueError as e:
+ if not loop_aware:
+ raise
+ raise HTTPException(status.HTTP_409_CONFLICT, str(e)) from e
# After the clear has succeeded, so a failed clear cannot take the task state with it.
# This is the only clear path that discards task state today; Dag-run clear and
@@ -1010,14 +1156,18 @@ def post_clear_task_instances(
if body.note is not None:
_patch_task_instance_note(
task_instance_body=body,
- tis=task_instances,
+ tis=list(task_instances),
user=user,
)
+ # The reload below refreshes with populate_existing, which would discard an unflushed note.
+ session.flush()
- task_instances = _reload_tis_with_rendered_fields(task_instances, session)
+ task_instances = _reload_tis_with_rendered_fields(list(task_instances), session)
return TaskInstanceCollectionResponse(
- task_instances=[TaskInstanceResponse.model_validate(ti) for ti in task_instances],
+ task_instances=[
+ task_coordinate_response(TaskInstanceResponse, ti, resolver) for ti in task_instances
+ ],
total_entries=len(task_instances),
)
@@ -1042,8 +1192,11 @@ def patch_task_group_instances(
session: SessionDep,
user: GetUserDep,
update_mask: list[str] | None = Query(None),
+ region_id: UUID | None = None,
+ region_index: int | None = None,
) -> TaskInstanceCollectionResponse:
"""Update the state of all task instances in a task group."""
+ body = patch_region_selection(body, region_id, region_index)
dag, tis, data = _patch_ti_group_validate_request(
dag_id, dag_run_id, group_id, dag_bag, body, session, update_mask
)
@@ -1065,12 +1218,14 @@ def patch_task_group_instances(
body=body,
data=data,
session=session,
+ selected=tis,
)
response_tis = _reload_tis_with_rendered_fields(response_tis, session)
+ resolver = TaskCoordinateResolver(dag_bag, session)
return TaskInstanceCollectionResponse(
- task_instances=[TaskInstanceResponse.model_validate(ti) for ti in response_tis],
+ task_instances=[task_coordinate_response(TaskInstanceResponse, ti, resolver) for ti in response_tis],
total_entries=len(response_tis),
)
@@ -1090,12 +1245,27 @@ def patch_task_group_instances_dry_run(
dag_bag: DagBagDep,
body: PatchTaskInstanceBody,
session: SessionDep,
+ region_id: UUID | None = None,
+ region_index: int | None = None,
) -> TaskInstanceCollectionResponse:
"""Dry-run of updating the state of all task instances in a task group."""
- dag = get_latest_version_of_dag(dag_bag, dag_id, session)
- tis = _get_task_group_task_instances(dag_id, dag_run_id, group_id, dag, session)
+ body = patch_region_selection(body, region_id, region_index)
+ dag, tis, data = _patch_ti_group_validate_request(
+ dag_id, dag_run_id, group_id, dag_bag, body, session, lock=False
+ )
- if body.new_state:
+ if body.new_state and body.region_id is not None:
+ tis = _patch_task_group_state(
+ group_id,
+ dag_run_id,
+ dag,
+ body,
+ data,
+ session=session,
+ selected=tis,
+ commit=False,
+ )
+ elif body.new_state:
tis = (
dag.set_task_group_state(
group_id=group_id,
@@ -1113,8 +1283,9 @@ def patch_task_group_instances_dry_run(
tis = _reload_tis_with_rendered_fields(tis, session)
+ resolver = TaskCoordinateResolver(dag_bag, session)
return TaskInstanceCollectionResponse(
- task_instances=[TaskInstanceResponse.model_validate(ti) for ti in tis],
+ task_instances=[task_coordinate_response(TaskInstanceResponse, ti, resolver) for ti in tis],
total_entries=len(tis),
)
@@ -1144,14 +1315,28 @@ def patch_task_instance_dry_run(
session: SessionDep,
map_index: int | None = None,
update_mask: list[str] | None = Query(None),
+ region_id: UUID | None = None,
+ region_index: int | None = None,
) -> TaskInstanceCollectionResponse:
"""Update a task instance dry_run mode."""
tis: Sequence[TI]
+ body = patch_region_selection(body, region_id, region_index)
dag, tis, data = _patch_ti_validate_request(
- dag_id, dag_run_id, task_id, dag_bag, body, session, map_index, update_mask
+ dag_id, dag_run_id, task_id, dag_bag, body, session, map_index, update_mask, lock=False
)
- if data.get("new_state"):
+ if data.get("new_state") and body.region_id is not None:
+ tis = _patch_task_instance_state(
+ task_id,
+ dag_run_id,
+ dag,
+ body,
+ data,
+ session,
+ selected=list(tis),
+ commit=False,
+ )
+ elif data.get("new_state"):
tis = (
dag.set_task_instance_state(
task_id=task_id,
@@ -1170,13 +1355,9 @@ def patch_task_instance_dry_run(
tis = _reload_tis_with_rendered_fields(tis, session)
+ resolver = TaskCoordinateResolver(dag_bag, session)
return TaskInstanceCollectionResponse(
- task_instances=[
- TaskInstanceResponse.model_validate(
- ti,
- )
- for ti in tis
- ],
+ task_instances=[task_coordinate_response(TaskInstanceResponse, ti, resolver) for ti in tis],
total_entries=len(tis),
)
@@ -1231,8 +1412,11 @@ def patch_task_instance(
session: SessionDep,
map_index: int | None = None,
update_mask: list[str] | None = Query(None),
+ region_id: UUID | None = None,
+ region_index: int | None = None,
) -> TaskInstanceCollectionResponse:
"""Update a task instance."""
+ body = patch_region_selection(body, region_id, region_index)
dag, tis, data = _patch_ti_validate_request(
dag_id, dag_run_id, task_id, dag_bag, body, session, map_index, update_mask
)
@@ -1250,6 +1434,8 @@ def patch_task_instance(
bulk_ti_body = BulkTaskInstanceBody(
task_id=task_id,
map_index=map_index,
+ region_id=body.region_id,
+ region_index=body.region_index,
new_state=body.new_state,
note=body.note,
include_upstream=body.include_upstream,
@@ -1257,24 +1443,22 @@ def patch_task_instance(
include_future=body.include_future,
include_past=body.include_past,
)
- _patch_task_instance_state(
+ updated_tis = _patch_task_instance_state(
task_id=task_id,
dag_run_id=dag_run_id,
dag=dag,
task_instance_body=bulk_ti_body,
data=data,
session=session,
+ selected=tis,
)
+ if body.region_id is not None:
+ tis = updated_tis
load_legacy_rendered_fields(tis, session=session)
-
+ resolver = TaskCoordinateResolver(dag_bag, session)
return TaskInstanceCollectionResponse(
- task_instances=[
- TaskInstanceResponse.model_validate(
- ti,
- )
- for ti in tis
- ],
+ task_instances=[task_coordinate_response(TaskInstanceResponse, ti, resolver) for ti in tis],
total_entries=len(tis),
)
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py
index 9ef6e15f3a0..fb3533af892 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py
@@ -54,13 +54,18 @@ from airflow.api_fastapi.core_api.datamodels.dag_run import (
ClearPartitionsBody,
DagRunMutableStates,
)
-from airflow.api_fastapi.core_api.datamodels.task_instances import NewTaskResponse
+from airflow.api_fastapi.core_api.datamodels.task_instances import NewTaskResponse, TaskInstanceResponse
from airflow.api_fastapi.core_api.services.public.common import BulkService
+from airflow.api_fastapi.core_api.services.public.task_coordinates import task_coordinate_response
from airflow.api_fastapi.core_api.services.public.task_instances import _emit_state_listener_hooks
from airflow.listeners.listener import get_listener_manager
from airflow.models.dagrun import DagRun, clear_partition_runs
from airflow.models.renderedtifields import load_legacy_rendered_fields
-from airflow.models.task_coordinates import enclosing_loop, public_map_index_expression
+from airflow.models.task_coordinates import (
+ TaskCoordinateResolver,
+ enclosing_loop,
+ public_map_index_expression,
+)
from airflow.models.taskinstance import TaskInstance
from airflow.models.xcom import XCOM_RETURN_KEY, XComModel, xcom_entity
from airflow.utils.session import create_session_async
@@ -132,7 +137,8 @@ def dry_run_clear_dag_run(
)
task_instances = list(session.scalars(ti_query))
load_legacy_rendered_fields(task_instances, session=session)
- return task_instances
+ resolver = TaskCoordinateResolver(dag_bag, session)
+ return [task_coordinate_response(TaskInstanceResponse, ti, resolver) for ti in task_instances]
def perform_clear_dag_run(
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_coordinates.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_coordinates.py
new file mode 100644
index 00000000000..2a155085a9f
--- /dev/null
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_coordinates.py
@@ -0,0 +1,113 @@
+# 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 dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, TypeVar
+from uuid import UUID
+
+from fastapi import HTTPException, status
+from pydantic import BaseModel
+
+from airflow._shared.state import TaskScope
+from airflow.models.dynamic_region import SENTINEL_REGION_ID, AmbiguousProducerError
+from airflow.models.task_coordinates import TaskCoordinateResolver
+
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
+
+ from airflow.models.dagbag import DBDagBag
+ from airflow.models.task_coordinates import TaskCoordinate
+
+
+Response = TypeVar("Response", bound=BaseModel)
+
+
+@dataclass
+class TaskCoordinateView:
+ """Present public mapping coordinates without changing ORM state."""
+
+ value: TaskCoordinate
+ resolver: TaskCoordinateResolver
+
+ def __getattr__(self, name: str) -> Any:
+ if name == "map_index":
+ return self.resolver.public_map_index(self.value)
+ return getattr(self.value, name)
+
+
+def task_coordinate_response(
+ schema: type[Response], value: TaskCoordinate, resolver: TaskCoordinateResolver
+) -> Response:
+ return schema.model_validate(TaskCoordinateView(value, resolver))
+
+
+def resolve_task_scope(
+ *,
+ dag_id: str,
+ run_id: str,
+ task_id: str,
+ session: Session,
+ dag_bag: DBDagBag,
+ map_index: int = -1,
+ region_id: UUID | None = None,
+ region_index: int | None = None,
+ all_map_indices: bool = False,
+) -> TaskScope:
+ """Resolve a public data address, including explicitly addressed retained data."""
+ if region_index is not None and region_id is None:
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, "region_index requires region_id")
+ if region_id is not None and region_index is not None:
+ return TaskScope(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ map_index=region_index,
+ region_id=region_id,
+ )
+ resolver = TaskCoordinateResolver(dag_bag, session)
+ if region_id == SENTINEL_REGION_ID or not resolver.has_regions(dag_id, run_id, task_id):
+ return TaskScope(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ map_index=map_index,
+ region_id=region_id or SENTINEL_REGION_ID,
+ )
+ try:
+ tasks = resolver.resolve(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ region_id=region_id,
+ map_indexes=None if all_map_indices else map_index,
+ )
+ 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 tasks:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, "Task instance not found for selected coordinates")
+ if len({ti.region_id for ti in tasks}) != 1 or (not all_map_indices and len(tasks) != 1):
+ raise HTTPException(status.HTTP_409_CONFLICT, "Select a region and index for this task instance")
+ return TaskScope(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ map_index=tasks[0].region_index,
+ region_id=tasks[0].region_id,
+ )
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py
index 4ad484c5a98..691cc5f254a 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py
@@ -18,17 +18,19 @@
from __future__ import annotations
from collections.abc import Sequence
-from typing import Literal
+from typing import Literal, TypeVar
+from uuid import UUID
import structlog
from fastapi import HTTPException, status
from fastapi.exceptions import RequestValidationError
from pydantic import ValidationError
-from sqlalchemy import select, tuple_
+from sqlalchemy import delete, select, tuple_
from sqlalchemy.orm import joinedload
from sqlalchemy.orm.session import Session
from airflow._shared.state import TaskScope
+from airflow.api.common.mark_tasks import get_run_ids
from airflow.api_fastapi.app import get_auth_manager
from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
from airflow.api_fastapi.common.dagbag import DagBagDep, get_latest_version_of_dag
@@ -48,16 +50,36 @@ from airflow.api_fastapi.core_api.datamodels.task_instances import (
)
from airflow.api_fastapi.core_api.security import GetUserDep
from airflow.api_fastapi.core_api.services.public.common import BulkService
+from airflow.api_fastapi.core_api.services.public.task_coordinates import resolve_task_scope
from airflow.configuration import conf
from airflow.listeners.listener import get_listener_manager
from airflow.models.dag import DagModel
+from airflow.models.dagbag import DBDagBag
+from airflow.models.dagrun import DagRun
+from airflow.models.dynamic_region import (
+ LOOP_DECISION_KEY,
+ SENTINEL_REGION_ID,
+ DynamicRegion,
+ load_region_ancestry,
+ loop_position,
+)
+from airflow.models.loop_clear import loop_coordinate_is_superseded, select_loop_clear_scope
from airflow.models.renderedtifields import load_legacy_rendered_fields
-from airflow.models.taskinstance import TaskInstance as TI
+from airflow.models.task_coordinates import (
+ TaskCoordinateResolver,
+ enclosing_loop,
+ public_map_index_expression,
+)
+from airflow.models.taskinstance import TaskInstance as TI, clear_task_instances
+from airflow.models.xcom import XComModelV2
from airflow.serialization.definitions.dag import SerializedDAG
from airflow.state.metastore import _get_db_backend
from airflow.utils.state import TaskInstanceState
log = structlog.get_logger(__name__)
+MutationAction = TypeVar(
+ "MutationAction", BulkUpdateAction[BulkTaskInstanceBody], BulkDeleteAction[BulkTaskInstanceBody]
+)
def _discard_task_state_store(tis: Sequence[TI], session: Session, *, event: str) -> None:
@@ -83,7 +105,8 @@ def _discard_task_state_store(tis: Sequence[TI], session: Session, *, event: str
dag_id=ti.dag_id,
run_id=ti.run_id,
task_id=ti.task_id,
- map_index=ti.map_index if ti.map_index is not None else -1,
+ region_id=ti.region_id,
+ region_index=ti.region_index,
)
try:
backend.clear(scope=scope, session=session)
@@ -111,7 +134,8 @@ def _clear_task_state_store_on_success(tis: Sequence[TI], session: Session) -> N
dag_id=ti.dag_id,
run_id=ti.run_id,
task_id=ti.task_id,
- map_index=ti.map_index if ti.map_index is not None else -1,
+ region_id=ti.region_id,
+ region_index=ti.region_index,
)
try:
backend.clear(scope=scope, session=session)
@@ -120,7 +144,7 @@ def _clear_task_state_store_on_success(tis: Sequence[TI], session: Session) -> N
dag_id=ti.dag_id,
run_id=ti.run_id,
task_id=ti.task_id,
- map_index=ti.map_index,
+ region_index=ti.region_index,
)
except Exception:
log.warning(
@@ -197,9 +221,14 @@ def _patch_ti_validate_request(
session: SessionDep,
map_index: int | None = -1,
update_mask: list[str] | None = None,
+ *,
+ lock: bool = True,
) -> tuple[SerializedDAG, list[TI], dict]:
+ _validate_region_selection(body)
dag = get_latest_version_of_dag(dag_bag, dag_id, session)
- if not dag.has_task(task_id):
+ if lock:
+ _lock_patch_runs(dag, dag_run_id, body, session)
+ if body.region_id is None and not dag.has_task(task_id):
raise HTTPException(status.HTTP_404_NOT_FOUND, f"Task '{task_id}' not found in Dag '{dag_id}'")
query = (
@@ -207,11 +236,36 @@ def _patch_ti_validate_request(
.where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id)
.options(joinedload(TI.rendered_task_instance_fields))
)
- if map_index is not None:
- query = query.where(TI.map_index == map_index)
- else:
- query = query.order_by(TI.map_index)
-
+ if body.region_id is None:
+ resolver = TaskCoordinateResolver(dag_bag, session)
+ for version in session.scalars(
+ select(TI.dag_version_id)
+ .where(
+ TI.dag_id == dag_id,
+ TI.run_id == dag_run_id,
+ TI.task_id == task_id,
+ )
+ .distinct()
+ ):
+ if enclosing_loop(resolver.get_task(dag_id, dag_run_id, task_id, dag_version_id=version)):
+ raise HTTPException(status.HTTP_409_CONFLICT, "Select a region and index for this loop task")
+ scope = resolve_task_scope(
+ dag_id=dag_id,
+ run_id=dag_run_id,
+ task_id=task_id,
+ session=session,
+ dag_bag=dag_bag,
+ map_index=map_index if map_index is not None else -1,
+ region_id=body.region_id,
+ region_index=body.region_index,
+ all_map_indices=map_index is None,
+ )
+ query = query.where(TI.region_id == scope.region_id)
+ if body.region_id is not None or map_index is not None:
+ query = query.where(TI.region_index == scope.region_index)
+ query = query.order_by(TI.region_index).execution_options(populate_existing=True)
+ if lock:
+ query = query.with_for_update(of=TI)
tis = session.scalars(query).all()
err_msg_404 = (
@@ -219,11 +273,55 @@ def _patch_ti_validate_request(
)
if len(tis) == 0:
raise HTTPException(status.HTTP_404_NOT_FOUND, err_msg_404)
+ if body.region_id is not None and tis[0].dag_version_id is not None:
+ pinned_dag = dag_bag.get_dag(tis[0].dag_version_id, session=session)
+ if pinned_dag is not None:
+ dag = pinned_dag
data = _validate_patch_task_instance_body(body, update_mask)
return dag, list(tis), data
+def _validate_region_selection(body: PatchTaskInstanceBody) -> None:
+ if (body.region_id is None) != (body.region_index is None):
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST, "region_id and region_index must be supplied together"
+ )
+ if body.region_id is not None and (body.include_past or body.include_future):
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, "Regional selection requires one explicit DagRun")
+
+
+def _lock_patch_runs(dag: SerializedDAG, run_id: str, body: PatchTaskInstanceBody, session: Session) -> None:
+ run_ids = get_run_ids(dag, run_id, body.include_future, body.include_past, session=session)
+ session.scalars(
+ select(DagRun.id)
+ .where(
+ DagRun.dag_id == dag.dag_id,
+ DagRun.run_id.in_(run_ids),
+ )
+ .order_by(DagRun.run_id)
+ .with_for_update()
+ ).all()
+
+
+def patch_region_selection(
+ body: PatchTaskInstanceBody, region_id: UUID | None, region_index: int | None
+) -> PatchTaskInstanceBody:
+ if region_id is None and region_index is None:
+ _validate_region_selection(body)
+ return body
+ if region_id is None or region_index is None:
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST, "region_id and region_index must be supplied together"
+ )
+ if body.region_id is not None or body.region_index is not None:
+ if (body.region_id, body.region_index) != (region_id, region_index):
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, "Body and query region coordinates conflict")
+ body = body.model_copy(update={"region_id": region_id, "region_index": region_index})
+ _validate_region_selection(body)
+ return body
+
+
def _get_task_group_task_ids(dag_id: str, task_group_id: str, dag: SerializedDAG) -> list[str]:
"""Return the ids of every task in a task group, resolved from the dag structure."""
task_group = dag.task_group_dict.get(task_group_id)
@@ -240,21 +338,68 @@ def _get_task_group_task_instances(
task_group_id: str,
dag: SerializedDAG,
session: Session,
+ body: PatchTaskInstanceBody | None = None,
+ dag_bag: DBDagBag | None = None,
) -> list[TI]:
"""Get all task instances in a task group for a specific DAG run."""
- task_ids = _get_task_group_task_ids(dag_id, task_group_id, dag)
-
query = (
select(TI)
.where(
TI.dag_id == dag_id,
TI.run_id == dag_run_id,
- TI.task_id.in_(task_ids),
)
- .order_by(TI.task_id, TI.map_index)
+ .order_by(TI.task_id, TI.region_id, TI.region_index)
+ .execution_options(populate_existing=True)
)
+ if body is None or body.region_id is None:
+ task_ids = _get_task_group_task_ids(dag_id, task_group_id, dag)
+ query = query.where(TI.task_id.in_(task_ids))
group_tis = list(session.scalars(query).all())
+ if body is not None:
+ _validate_region_selection(body)
+ if body.region_id is None:
+ resolver = TaskCoordinateResolver(dag_bag or DBDagBag(), session)
+ if any(
+ ti.region_id != SENTINEL_REGION_ID
+ and enclosing_loop(
+ resolver.get_task(dag_id, dag_run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+ )
+ for ti in group_tis
+ ):
+ raise HTTPException(status.HTTP_409_CONFLICT, "Select a region and index for this loop group")
+ elif body.region_index is not None:
+ region = session.get(DynamicRegion, body.region_id)
+ if region is None or (region.dag_id, region.run_id) != (dag_id, dag_run_id):
+ raise HTTPException(status.HTTP_404_NOT_FOUND, "Selected loop region not found")
+ regions = load_region_ancestry(
+ {ti.region_id for ti in group_tis} | {body.region_id},
+ dag_id=dag_id,
+ run_id=dag_run_id,
+ session=session,
+ )
+ position = loop_position(regions, body.region_id, body.region_index, region.node_id)
+ group_tis = [
+ ti
+ for ti in group_tis
+ if loop_position(regions, ti.region_id, ti.region_index, region.node_id) == position
+ ]
+ resolver = TaskCoordinateResolver(dag_bag or DBDagBag(), session)
+ members = []
+ for ti in group_tis:
+ task = resolver.get_task(dag_id, dag_run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+ loop = enclosing_loop(task)
+ if loop is None or loop.node_id != region.node_id:
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST, "Group coordinates must identify an enclosing loop pass"
+ )
+ group = task.task_group
+ while group is not None:
+ if group.group_id == task_group_id:
+ members.append(ti)
+ break
+ group = group.parent_group
+ group_tis = members
if not group_tis:
raise HTTPException(
status.HTTP_404_NOT_FOUND,
@@ -272,10 +417,23 @@ def _patch_ti_group_validate_request(
body: PatchTaskInstanceBody,
session: SessionDep,
update_mask: list[str] | None = None,
+ *,
+ lock: bool = True,
) -> tuple[SerializedDAG, list[TI], dict]:
"""Validate and prepare data for task group patch request."""
dag = get_latest_version_of_dag(dag_bag, dag_id, session)
- tis = _get_task_group_task_instances(dag_id, dag_run_id, task_group_id, dag, session)
+ if lock:
+ _lock_patch_runs(dag, dag_run_id, body, session)
+ tis = _get_task_group_task_instances(dag_id, dag_run_id, task_group_id, dag, session, body, dag_bag)
+ if lock:
+ tis = list(
+ session.scalars(
+ select(TI)
+ .where(TI.id.in_([ti.id for ti in tis]))
+ .with_for_update(of=TI)
+ .execution_options(populate_existing=True)
+ )
+ )
data = _validate_patch_task_instance_body(body, update_mask)
return dag, tis, data
@@ -288,7 +446,23 @@ def _patch_task_instance_state(
task_instance_body: BulkTaskInstanceBody | PatchTaskInstanceBody,
data: dict,
session: Session,
+ selected: list[TI] | None = None,
+ commit: bool = True,
) -> list[TI]:
+ if task_instance_body.region_id is not None:
+ if selected is None:
+ selected = list(
+ session.scalars(
+ select(TI).where(
+ TI.dag_id == dag.dag_id,
+ TI.run_id == dag_run_id,
+ TI.task_id == task_id,
+ TI.region_id == task_instance_body.region_id,
+ TI.region_index == task_instance_body.region_index,
+ )
+ )
+ )
+ return _patch_selected_task_state(selected, task_instance_body, data, session=session, commit=commit)
map_index = getattr(task_instance_body, "map_index", None)
map_indexes = None if map_index is None else [map_index]
@@ -318,6 +492,75 @@ def _patch_task_instance_state(
return updated_tis
+def _patch_selected_task_state(
+ selected: list[TI], body: PatchTaskInstanceBody, data: dict, *, session: Session, commit: bool
+) -> list[TI]:
+ if commit and selected:
+ session.scalars(
+ select(DagRun.id)
+ .where(DagRun.dag_id == selected[0].dag_id, DagRun.run_id == selected[0].run_id)
+ .with_for_update()
+ ).all()
+ scope = select_loop_clear_scope(
+ selected,
+ upstream=body.include_upstream,
+ downstream=body.include_downstream,
+ later_loop_iterations=False,
+ session=session,
+ )
+ query = select(TI).where(TI.id.in_(scope.retry_ids))
+ if commit:
+ query = query.with_for_update(of=TI).execution_options(populate_existing=True)
+ tis = list(session.scalars(query))
+ changed = [ti for ti in tis if ti.state != data["new_state"]]
+ if not commit:
+ return changed
+ if not changed:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT, f"Selected task instances are already in {data['new_state']} state"
+ )
+ regions = {
+ region.id: region
+ for region in session.scalars(
+ select(DynamicRegion).where(
+ DynamicRegion.dag_id == tis[0].dag_id,
+ DynamicRegion.run_id == tis[0].run_id,
+ )
+ )
+ }
+ superseded = {
+ ti.id for ti in tis if loop_coordinate_is_superseded(ti.region_id, ti.region_index, regions)
+ }
+ downstream = select_loop_clear_scope(
+ [ti for ti in tis if ti.id not in superseded], later_loop_iterations=False, session=session
+ )
+ for ti in changed:
+ if ti.operator == "LoopGateOperator":
+ session.execute(
+ delete(XComModelV2).where(
+ XComModelV2.task_instance_id == ti.id, XComModelV2.key == LOOP_DECISION_KEY
+ )
+ )
+ if ti.state == TaskInstanceState.RESTARTING and ti.id in superseded:
+ ti.complete_restart(session=session, terminal_outcome=data["new_state"])
+ else:
+ ti.set_state(data["new_state"], session=session)
+ failures = list(
+ session.scalars(
+ select(TI).where(
+ TI.id.in_(downstream.retry_ids - scope.retry_ids),
+ TI.state.in_((TaskInstanceState.FAILED, TaskInstanceState.UPSTREAM_FAILED)),
+ )
+ )
+ )
+ if failures:
+ clear_task_instances(failures, session=session)
+ if data["new_state"] == TaskInstanceState.SUCCESS:
+ _clear_task_state_store_on_success(changed, session)
+ _emit_state_listener_hooks(changed, data["new_state"])
+ return changed
+
+
def _patch_task_group_state(
group_id: str,
dag_run_id: str,
@@ -326,8 +569,14 @@ def _patch_task_group_state(
data: dict,
*,
session: Session,
+ selected: list[TI] | None = None,
+ commit: bool = True,
) -> list[TI]:
"""Update the state of all task instances in a task group."""
+ if body.region_id is not None:
+ if selected is None:
+ selected = _get_task_group_task_instances(dag.dag_id, dag_run_id, group_id, dag, session, body)
+ return _patch_selected_task_state(selected, body, data, session=session, commit=commit)
updated_tis = dag.set_task_group_state(
group_id=group_id,
run_id=dag_run_id,
@@ -484,17 +733,35 @@ class BulkTaskInstanceService(BulkService[BulkTaskInstanceBody]):
# Filter at database level using exact tuple matching instead of fetching all combinations
# and filtering in Python
task_keys_list = list(task_keys)
- query = select(TI).where(tuple_(TI.dag_id, TI.run_id, TI.task_id, TI.map_index).in_(task_keys_list))
-
- task_instances = self.session.scalars(query).all()
- task_instances_map = {
- (ti.dag_id, ti.run_id, ti.task_id, ti.map_index if ti.map_index is not None else -1): ti
- for ti in task_instances
- }
+ public_index = public_map_index_expression(TI)
+ query = select(TI, public_index).where(
+ tuple_(TI.dag_id, TI.run_id, TI.task_id, public_index).in_(task_keys_list)
+ )
+ rows = self.session.execute(query).all()
+ self._reject_unscoped_loop_tasks([ti for ti, _ in rows])
+ task_instances_map = {}
+ for ti, index in rows:
+ key = (ti.dag_id, ti.run_id, ti.task_id, index)
+ if key in task_instances_map:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT, "Select region coordinates for loop task instances"
+ )
+ task_instances_map[key] = ti
matched_task_keys = set(task_instances_map.keys())
not_found_task_keys = task_keys - matched_task_keys
return task_instances_map, matched_task_keys, not_found_task_keys
+ def _reject_unscoped_loop_tasks(self, tis: Sequence[TI]) -> None:
+ resolver = TaskCoordinateResolver(self.dag_bag, self.session)
+ for ti in tis:
+ if ti.region_id == SENTINEL_REGION_ID:
+ continue
+ task = resolver.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+ if enclosing_loop(task) is not None:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT, "Select region coordinates for loop task instances"
+ )
+
def _perform_update(
self,
entity: BulkTaskInstanceBody,
@@ -545,13 +812,92 @@ class BulkTaskInstanceService(BulkService[BulkTaskInstanceBody]):
}
)
+ def _handle_regional_bulk(
+ self, action: MutationAction, results: BulkActionResponse
+ ) -> tuple[MutationAction, set[tuple[str, str, str, int]], set[tuple[str, str, str]]]:
+ regional = [
+ entity
+ for entity in action.entities
+ if isinstance(entity, BulkTaskInstanceBody)
+ and (entity.region_id is not None or entity.region_index is not None)
+ ]
+ deleting = isinstance(action, BulkDeleteAction)
+ specific, whole = self._categorize_entities(
+ action.entities, results, method="DELETE" if deleting else "PUT", action_name=action.action.value
+ )
+ keys = {key[:3] for key in specific} | whole
+ run_keys = {key[:2] for key in keys}
+ for entity in action.entities:
+ if isinstance(entity, BulkTaskInstanceBody) and (entity.include_future or entity.include_past):
+ dag_id, run_id, task_id, _ = self._extract_task_identifiers(entity)
+ if (dag_id, run_id, task_id) in keys:
+ dag = get_latest_version_of_dag(self.dag_bag, dag_id, self.session)
+ run_keys.update(
+ (dag_id, selected_run)
+ for selected_run in get_run_ids(
+ dag, run_id, entity.include_future, entity.include_past, session=self.session
+ )
+ )
+ self.session.scalars(
+ select(DagRun)
+ .where(
+ tuple_(DagRun.dag_id, DagRun.run_id).in_(run_keys),
+ )
+ .order_by(DagRun.dag_id, DagRun.run_id)
+ .with_for_update()
+ ).all()
+ for entity in regional:
+ dag_id, run_id, task_id, map_index = self._extract_task_identifiers(entity)
+ if (dag_id, run_id, task_id) not in keys:
+ continue
+ try:
+ dag, tis, data = _patch_ti_validate_request(
+ dag_id,
+ run_id,
+ task_id,
+ self.dag_bag,
+ entity,
+ self.session,
+ map_index,
+ getattr(action, "update_mask", None),
+ )
+ if deleting:
+ for ti in tis:
+ self.session.delete(ti)
+ else:
+ if "note" in data:
+ _patch_task_instance_note(entity, tis, self.user)
+ if "new_state" in data:
+ _patch_task_instance_state(
+ task_id, run_id, dag, entity, data, self.session, selected=tis
+ )
+ results.success.extend(
+ str(ti.id)
+ if ti.region_id != SENTINEL_REGION_ID
+ else f"{dag_id}.{run_id}.{task_id}[{ti.region_index}]"
+ for ti in tis
+ )
+ except HTTPException as error:
+ if (
+ error.status_code == status.HTTP_404_NOT_FOUND
+ and action.action_on_non_existence != BulkActionNotOnExistence.FAIL
+ ):
+ continue
+ results.errors.append({"error": str(error.detail), "status_code": error.status_code})
+ remaining = [entity for entity in action.entities if entity not in regional]
+ remaining_keys = {self._extract_task_identifiers(entity) for entity in remaining}
+ return (
+ action.model_copy(update={"entities": remaining}),
+ specific & remaining_keys,
+ whole & {key[:3] for key in remaining_keys if key[3] is None},
+ )
+
def handle_bulk_update(
self, action: BulkUpdateAction[BulkTaskInstanceBody], results: BulkActionResponse
) -> None:
"""Bulk Update Task Instances."""
- # Validate and categorize entities into specific and all map index update sets
- update_specific_map_index_task_keys, update_all_map_index_task_keys = self._categorize_entities(
- action.entities, results, method="PUT", action_name=action.action.value
+ action, update_specific_map_index_task_keys, update_all_map_index_task_keys = (
+ self._handle_regional_bulk(action, results)
)
try:
@@ -611,6 +957,7 @@ class BulkTaskInstanceService(BulkService[BulkTaskInstanceBody]):
).all()
# Group task instances by (dag_id, run_id, task_id)
+ self._reject_unscoped_loop_tasks(batch_task_instances)
task_instances_by_key: dict[tuple[str, str, str], list[TI]] = {}
for ti in batch_task_instances:
key = (ti.dag_id, ti.run_id, ti.task_id)
@@ -636,7 +983,7 @@ class BulkTaskInstanceService(BulkService[BulkTaskInstanceBody]):
dag_id=dag_id,
dag_run_id=run_id,
task_id=task_id,
- map_index=ti.map_index if ti.map_index is not None else -1,
+ map_index=ti.region_index,
entity=entity,
results=results,
update_mask=action.update_mask,
@@ -651,9 +998,8 @@ class BulkTaskInstanceService(BulkService[BulkTaskInstanceBody]):
self, action: BulkDeleteAction[BulkTaskInstanceBody], results: BulkActionResponse
) -> None:
"""Bulk delete task instances."""
- # Validate and categorize entities into specific and all map index delete sets
- delete_specific_map_index_task_keys, delete_all_map_index_task_keys = self._categorize_entities(
- action.entities, results, method="DELETE", action_name=action.action.value
+ action, delete_specific_map_index_task_keys, delete_all_map_index_task_keys = (
+ self._handle_regional_bulk(action, results)
)
try:
@@ -699,6 +1045,7 @@ class BulkTaskInstanceService(BulkService[BulkTaskInstanceBody]):
).all()
# Group task instances by (dag_id, run_id, task_id) for efficient lookup
+ self._reject_unscoped_loop_tasks(batch_task_instances)
task_instances_by_key: dict[tuple[str, str, str], list[TI]] = {}
for ti in batch_task_instances:
key = (ti.dag_id, ti.run_id, ti.task_id)
@@ -721,7 +1068,7 @@ class BulkTaskInstanceService(BulkService[BulkTaskInstanceBody]):
dag_id=dag_id, run_id=run_id, task_id=task_id, session=self.session
)
for ti in all_task_instances:
- results.success.append(f"{dag_id}.{run_id}.{task_id}[{ti.map_index}]")
+ results.success.append(f"{dag_id}.{run_id}.{task_id}[{ti.region_index}]")
except HTTPException as e:
results.errors.append({"error": f"{e.detail}", "status_code": e.status_code})
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 43fb7660357..28bdcdded1a 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 import trace
from opentelemetry.trace import StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from pydantic import JsonValue, ValidationError
-from sqlalchemy import and_, func, or_, tuple_, union, update
+from sqlalchemy import and_, exists, func, or_, tuple_, union, update
from sqlalchemy.engine import CursorResult
from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError
from sqlalchemy.orm import contains_eager, joinedload
@@ -417,6 +417,23 @@ def ti_update_state(
if isinstance(ti_patch_payload, TITerminalStatePayload) and (
ti_patch_payload.state == TerminalStateNonSuccess.SERVER_TERMINATED
):
+ in_region = session.scalar(
+ select(or_(TI.region_id != SENTINEL_REGION_ID, TI.region_index >= 0)).where(
+ TI.id == task_instance_id
+ )
+ )
+ if in_region:
+ session.execute(
+ select(DR.id)
+ .where(
+ exists().where(
+ TI.id == task_instance_id,
+ TI.dag_id == DR.dag_id,
+ TI.run_id == DR.run_id,
+ )
+ )
+ .with_for_update()
+ ).all()
ti = session.scalar(
select(TI)
.where(TI.id == task_instance_id)
diff --git a/airflow-core/src/airflow/cli/commands/dag_command.py b/airflow-core/src/airflow/cli/commands/dag_command.py
index 1f96179bfbb..a05d501d3d3 100644
--- a/airflow-core/src/airflow/cli/commands/dag_command.py
+++ b/airflow-core/src/airflow/cli/commands/dag_command.py
@@ -44,8 +44,8 @@ from airflow.exceptions import AirflowConfigException, AirflowException
from airflow.jobs.job import Job
from airflow.models import DagModel, DagRun, TaskInstance
from airflow.models.errors import ParseImportError
+from airflow.models.loop_clear import clear_task_instances_for_runs
from airflow.models.serialized_dag import SerializedDagModel
-from airflow.models.taskinstance import clear_task_instances
from airflow.timetables.base import TimeRestriction
from airflow.utils import cli as cli_utils
from airflow.utils.cli import (
@@ -234,7 +234,7 @@ def _bulk_clear_runs(
tis = session.scalars(ti_query).all()
if not tis:
continue
- clear_task_instances(list(tis), session=session)
+ clear_task_instances_for_runs(tis, session=session)
session.flush()
cleared += len(tis)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 189e1e901fa..510b61b5c40 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -1542,10 +1542,36 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
return len(event_buffer)
# Check state of finished tasks
+ candidates = session.execute(
+ select(TI.id, TI.dag_id, TI.run_id).where(
+ TI.working_set.is_(True), TI.id.in_([key.id for key in tis_with_right_state])
+ )
+ ).all()
+ run_query = (
+ select(DagRun.dag_id, DagRun.run_id)
+ .where(
+ tuple_(DagRun.dag_id, DagRun.run_id).in_(
+ {(dag_id, run_id) for _, dag_id, run_id in candidates}
+ )
+ )
+ .order_by(DagRun.dag_id, DagRun.run_id)
+ )
+ locked_runs = set(
+ session.execute(with_row_locks(run_query, of=DagRun, session=session, skip_locked=True))
+ )
+ available_ids = []
+ num_requeued = 0
+ for ti_id, dag_id, run_id in candidates:
+ if (dag_id, run_id) in locked_runs:
+ available_ids.append(ti_id)
+ else:
+ key = TaskInstanceUuid(ti_id)
+ executor.event_buffer.setdefault(key, event_buffer.pop(key))
+ num_requeued += 1
asset_loader, alias_loader = _eager_load_dag_run_for_validation()
query = (
select(TI)
- .where(TI.id.in_([key.id for key in tis_with_right_state]))
+ .where(TI.id.in_(available_ids))
.options(selectinload(TI.dag_model))
.options(asset_loader)
.options(alias_loader)
@@ -1771,15 +1797,28 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
# Update task state - emails are handled by DAG processor now
ti.handle_failure(error=msg, session=session)
- for task_id in tis_with_right_state:
- if task_id in event_buffer:
- cls.logger().warning(
- "Discarding executor event for task instance %s (coordinates=%s): no matching task instance was "
- "returned; it may no longer exist or may be locked by another scheduler",
- task_id,
- event_coordinates.get(task_id),
+ unprocessed_keys = [key for key in tis_with_right_state if key in event_buffer]
+ if unprocessed_keys:
+ # The lock query skips rows another transaction holds; only rows that are still live are retried.
+ live_ids = set(
+ session.scalars(
+ select(TI.id).where(
+ TI.working_set.is_(True), TI.id.in_([key.id for key in unprocessed_keys])
+ )
)
- cls._emit_executor_events_batch_metrics(num_events)
+ )
+ for key in unprocessed_keys:
+ if key.id in live_ids:
+ executor.event_buffer.setdefault(key, event_buffer.pop(key))
+ num_requeued += 1
+ else:
+ cls.logger().warning(
+ "Discarding executor event for task instance %s (coordinates=%s): the task instance "
+ "no longer exists or is no longer current",
+ key,
+ event_coordinates.get(key),
+ )
+ cls._emit_executor_events_batch_metrics(num_events - num_requeued)
return len(event_buffer)
def _execute(self) -> int | None:
@@ -3891,6 +3930,19 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
)
.where(TI.queued_by_job_id == self.job.id)
)
+ run_query = (
+ select(DagRun.dag_id, DagRun.run_id)
+ .where(
+ tuple_(DagRun.dag_id, DagRun.run_id).in_(
+ query.with_only_columns(TI.dag_id, TI.run_id).distinct()
+ )
+ )
+ .order_by(DagRun.dag_id, DagRun.run_id)
+ )
+ locked_runs = list(
+ session.execute(with_row_locks(run_query, of=DagRun, session=session, skip_locked=True))
+ )
+ query = query.where(tuple_(TI.dag_id, TI.run_id).in_(locked_runs))
# Lock the rows (FOR UPDATE, of=TI so the FOR UPDATE isn't applied to the joined dag_model)
# so a worker can't commit a terminal state on the same TI between this scan and the
# handle_failure() in the purge that follows in the same transaction. skip_locked keeps HA
diff --git a/airflow-core/src/airflow/models/dagbag.py b/airflow-core/src/airflow/models/dagbag.py
index 4415d8b0b2b..1665699bec1 100644
--- a/airflow-core/src/airflow/models/dagbag.py
+++ b/airflow-core/src/airflow/models/dagbag.py
@@ -22,7 +22,6 @@ import math
import time
from collections.abc import MutableMapping
from contextlib import nullcontext
-from functools import cache
from threading import RLock
from typing import TYPE_CHECKING, Any, NamedTuple
from uuid import UUID
@@ -46,9 +45,6 @@ if TYPE_CHECKING:
from airflow.serialization.definitions.dag import SerializedDAG
-SHARED_DAG_CACHE_SIZE = 64
-
-
class _CacheEntry(NamedTuple):
"""A cached deserialized DAG plus the metadata needed to detect staleness on lookup."""
@@ -293,23 +289,6 @@ class CachedDBDagBag(DBDagBag):
stats.gauge(f"{self._stats_prefix}.cache_size", size, rate=rate)
-@cache
-def get_shared_dag_bag() -> DBDagBag:
- """
- Return the process-wide Dag cache for structural lookups on hot paths.
-
- Scheduling, log addressing and task coordinate resolution only need a Dag's shape, and a
- version's definition is immutable, so one revalidated cache is reused rather than
- deserializing the serialized Dag for every pass or request.
- """
- return CachedDBDagBag(
- load_op_links=False,
- cache_size=SHARED_DAG_CACHE_SIZE,
- cache_ttl=0,
- stats_prefix="shared.dag_bag",
- )
-
-
def generate_md5_hash(context):
bundle_name = context.get_current_parameters()["bundle_name"]
relative_fileloc = context.get_current_parameters()["relative_fileloc"]
diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py
index 03d3e9bebb0..f4c6cf5619b 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -94,7 +94,7 @@ from airflow.models.dynamic_region import (
DynamicRegion,
public_region_filter,
)
-from airflow.models.taskinstance import TaskInstance as TI, _add_and_prime_mapped_ti, clear_task_instances
+from airflow.models.taskinstance import TaskInstance as TI, _add_and_prime_mapped_ti
from airflow.models.tasklog import LogTemplate
from airflow.serialization.definitions.deadline import SerializedReferenceModels
from airflow.serialization.definitions.notset import NOTSET, ArgNotSet, is_arg_set
@@ -1468,6 +1468,8 @@ class DagRun(Base, LoggingMixin):
@provide_session
def task_instance_scheduling_decisions(self, *, session: Session = NEW_SESSION) -> TISchedulingDecision:
tis = self.get_task_instances(session=session, state=State.task_states)
+ reconciled = self._reconcile_legacy_expansions(session=session)
+ tis = list({ti.id: ti for ti in (*tis, *reconciled)}.values())
self.log.debug("number of tis tasks for %s: %s task(s)", self, len(tis))
def _filter_tis_and_exclude_removed(dag: SerializedDAG, tis: list[TI]) -> Iterable[TI]:
@@ -1901,6 +1903,97 @@ class DagRun(Base, LoggingMixin):
tags={**self.stats_tags, "dag_id": self.dag_id},
)
+ def _reconcile_legacy_expansions(
+ self, *, session: Session, dag_version_id: UUID | None = None
+ ) -> list[TI]:
+ from airflow.models.task_coordinates import TaskCoordinateResolver
+
+ pending = (
+ select(DynamicRegion)
+ .where(
+ DynamicRegion.dag_id == self.dag_id,
+ DynamicRegion.run_id == self.run_id,
+ DynamicRegion.parent_region_id.is_(None),
+ DynamicRegion.forked_from_region_id.is_(None),
+ ~select(TI.id)
+ .where(
+ TI.working_set.is_(True),
+ TI.dag_id == DynamicRegion.dag_id,
+ TI.run_id == DynamicRegion.run_id,
+ TI.region_id == DynamicRegion.id,
+ )
+ .exists(),
+ ~select(TI.id)
+ .where(
+ TI.working_set.is_(True),
+ TI.dag_id == DynamicRegion.dag_id,
+ TI.run_id == DynamicRegion.run_id,
+ TI.task_id == DynamicRegion.node_id,
+ TI.region_id == SENTINEL_REGION_ID,
+ )
+ .exists(),
+ select(TI.id)
+ .where(
+ TI.working_set.is_(None),
+ TI.dag_id == DynamicRegion.dag_id,
+ TI.run_id == DynamicRegion.run_id,
+ TI.task_id == DynamicRegion.node_id,
+ TI.region_id == SENTINEL_REGION_ID,
+ TI.archived_reason == "superseded",
+ )
+ .exists(),
+ )
+ .execution_options(include_all_attempts=True)
+ )
+ regions = session.scalars(pending).all()
+ if not regions:
+ return []
+ session.execute(
+ select(DagRun)
+ .where(DagRun.id == self.id)
+ .with_for_update()
+ .execution_options(populate_existing=True)
+ ).scalar_one()
+ version_id = dag_version_id or self.created_dag_version_id
+ if version_id is None:
+ return []
+ dag = TaskCoordinateResolver.for_dag(self.dag, session).get_dag(version_id)
+ if dag is None:
+ return []
+ live = list(
+ session.scalars(
+ select(TI)
+ .where(
+ TI.working_set.is_(True),
+ TI.dag_id == self.dag_id,
+ TI.run_id == self.run_id,
+ or_(
+ TI.region_id.in_([region.id for region in regions]),
+ and_(
+ TI.region_id == SENTINEL_REGION_ID,
+ TI.task_id.in_([region.node_id for region in regions]),
+ ),
+ ),
+ )
+ .execution_options(populate_existing=True)
+ ).all()
+ )
+ occupied_regions = {ti.region_id for ti in live}
+ archiving_tasks = {ti.task_id for ti in live if ti.region_id == SENTINEL_REGION_ID}
+ for region in regions:
+ if region.id in occupied_regions or region.node_id in archiving_tasks:
+ continue
+ if not dag.has_task(region.node_id):
+ continue
+ task = dag.get_task(region.node_id)
+ if not task.get_needs_expansion():
+ continue
+ ti = TI(task, run_id=self.run_id, dag_version_id=version_id, region_id=region.id)
+ _add_and_prime_mapped_ti(ti, task, self, session=session)
+ live.append(ti)
+ session.flush()
+ return live
+
@provide_session
def verify_integrity(self, *, session: Session = NEW_SESSION, dag_version_id: UUID) -> None:
"""
@@ -1918,6 +2011,7 @@ class DagRun(Base, LoggingMixin):
hook_is_noop: Literal[True, False] = getattr(task_instance_mutation_hook, "is_noop", False)
dag = self.get_dag()
+ self._reconcile_legacy_expansions(session=session, dag_version_id=dag_version_id)
task_ids = self._check_for_removed_or_restored_tasks(
dag, task_instance_mutation_hook, session=session
)
@@ -2150,24 +2244,20 @@ class DagRun(Base, LoggingMixin):
session: Session,
) -> None:
"""Consume a gate decision while the caller holds the DagRun and TI locks."""
+ from airflow.models.loop_clear import (
+ loop_gate_has_later_pass,
+ loop_gate_waits_for_archival,
+ loop_successor_region,
+ )
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")
+ if state == TaskInstanceState.SUCCESS and loop_gate_waits_for_archival(gate, session=session):
+ raise InvalidLoopDecision("Loop gate must wait for superseded executions to terminate")
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)
- )
+ later_gate = loop_gate_has_later_pass(gate, group, session=session)
if state != TaskInstanceState.SUCCESS or later_gate:
if signal is not None:
session.delete(signal)
@@ -2193,7 +2283,7 @@ class DagRun(Base, LoggingMixin):
group.iter_tasks(),
creator,
session=session,
- parent_region=(gate.region_id, gate.region_index + 1),
+ parent_region=(loop_successor_region(gate, session=session), gate.region_index + 1),
)
)
self._create_task_instances(
@@ -2207,6 +2297,7 @@ class DagRun(Base, LoggingMixin):
*,
session: Session,
) -> CreatedTasks:
+ from airflow.models.loop_clear import loop_coordinate_is_superseded
from airflow.models.task_coordinates import enclosing_loop
ordinary_tasks = []
@@ -2219,6 +2310,18 @@ class DagRun(Base, LoggingMixin):
else:
ordinary_tasks.append(task)
yield from self._create_tasks(ordinary_tasks, task_creator, session=session, expand_literals=True)
+ regions = (
+ {
+ region.id: region
+ for region in session.scalars(
+ select(DynamicRegion).where(
+ DynamicRegion.dag_id == self.dag_id, DynamicRegion.run_id == self.run_id
+ )
+ )
+ }
+ if loop_tasks
+ else {}
+ )
for group_id, members in loop_tasks.items():
coordinates = (
session.execute(
@@ -2248,6 +2351,8 @@ class DagRun(Base, LoggingMixin):
session.flush()
coordinates = [(region.id, 0)]
for region_id, region_index in coordinates:
+ if loop_coordinate_is_superseded(region_id, region_index, regions):
+ continue
yield from self._create_tasks(
members,
task_creator,
@@ -2712,6 +2817,8 @@ def clear_partition_runs(
:meta private:
"""
+ from airflow.models.loop_clear import clear_task_instances_for_runs
+
stmt = select(DagRun).where(DagRun.dag_id == dag_id)
if run_id is not None:
stmt = stmt.where(DagRun.run_id == run_id)
@@ -2732,6 +2839,15 @@ def clear_partition_runs(
)
stmt = stmt.order_by(DagRun.partition_date, DagRun.run_id)
+ if clear_tis and not dry_run:
+ session.execute(
+ stmt.with_only_columns(DagRun.id)
+ .where(select(TI.id).where(TI.dag_id == DagRun.dag_id, TI.run_id == DagRun.run_id).exists())
+ .order_by(None)
+ .order_by(DagRun.dag_id, DagRun.run_id)
+ .with_for_update()
+ ).all()
+
dag_runs_cleared = 0
ti_buffer_run_ids: list[str] = []
ti_carry: list[TI] = []
@@ -2754,10 +2870,18 @@ def clear_partition_runs(
while len(ti_carry) >= _TI_CHUNK_SIZE:
slice_tis = ti_carry[:_TI_CHUNK_SIZE]
del ti_carry[:_TI_CHUNK_SIZE]
- clear_task_instances(slice_tis, session=session)
+ clear_task_instances_for_runs(
+ slice_tis,
+ session=session,
+ whole_task_keys={(ti.dag_id, ti.run_id, ti.task_id) for ti in slice_tis},
+ )
flushed += len(slice_tis)
if drain and ti_carry:
- clear_task_instances(ti_carry, session=session)
+ clear_task_instances_for_runs(
+ ti_carry,
+ session=session,
+ whole_task_keys={(ti.dag_id, ti.run_id, ti.task_id) for ti in ti_carry},
+ )
flushed += len(ti_carry)
return flushed
diff --git a/airflow-core/src/airflow/models/loop_clear.py b/airflow-core/src/airflow/models/loop_clear.py
new file mode 100644
index 00000000000..c8a3c828969
--- /dev/null
+++ b/airflow-core/src/airflow/models/loop_clear.py
@@ -0,0 +1,449 @@
+# 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 collections.abc import Collection, Iterator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal
+from uuid import UUID
+
+from sqlalchemy import exists, select, tuple_
+
+from airflow.exceptions import AirflowClearRunningTaskException
+from airflow.models.dagbag import DBDagBag
+from airflow.models.dynamic_region import (
+ SENTINEL_REGION_ID,
+ DynamicRegion,
+ load_region_ancestry,
+ loop_position,
+)
+from airflow.models.task_coordinates import TaskCoordinateResolver, enclosing_loop
+from airflow.models.taskinstance import TaskInstance, _get_relevant_map_indexes, clear_task_instances
+from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
+from airflow.utils.state import DagRunState, TaskInstanceState
+
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
+
+ from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup
+
+
+@dataclass(frozen=True)
+class LoopClearScope:
+ """Live executions to retry and generated executions to archive."""
+
+ retry_ids: frozenset[UUID]
+ archive_ids: frozenset[UUID]
+
+
+def select_loop_clear_scope(
+ selected: Collection[TaskInstance],
+ *,
+ whole_expansion_ids: Collection[UUID] = (),
+ upstream: bool = False,
+ downstream: bool = True,
+ later_loop_iterations: bool = True,
+ session: Session,
+) -> LoopClearScope:
+ """Select loop clear executions; mutation callers must hold the DagRun lock."""
+ if not selected:
+ return LoopClearScope(frozenset(), frozenset())
+ runs = {(ti.dag_id, ti.run_id) for ti in selected}
+ if len(runs) != 1:
+ raise ValueError("Loop clear selection must belong to one DagRun")
+ dag_id, run_id = runs.pop()
+ live = {
+ ti.id: ti
+ for ti in session.scalars(
+ select(TaskInstance)
+ .where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.working_set.is_(True),
+ )
+ .execution_options(populate_existing=True)
+ )
+ }
+ selected_ids = {ti.id for ti in selected}
+ if not selected_ids <= live.keys():
+ raise ValueError("Loop clear selection contains executions that are no longer live")
+ if not set(whole_expansion_ids) <= selected_ids:
+ raise ValueError("Whole-task selection requires an explicitly selected execution")
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+ regions = load_region_ancestry(
+ {ti.region_id for ti in live.values()}, dag_id=dag_id, run_id=run_id, session=session
+ )
+ retry = {ti_id: live[ti_id] for ti_id in selected_ids}
+ for ti in tuple(retry.values()):
+ task = resolver.get_task(dag_id, run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+ group = enclosing_loop(task)
+ if group is not None and loop_position(regions, ti.region_id, ti.region_index, group.node_id) is None:
+ raise ValueError("Loop clear selection requires an execution inside its pinned loop")
+ if ti.id in whole_expansion_ids:
+ if not task.get_needs_expansion():
+ raise ValueError("Whole-expansion selection requires a mapped task")
+ retry.update(
+ (other.id, other)
+ for other in resolver.resolve(dag_id=dag_id, run_id=run_id, task_id=ti.task_id, caller=ti)
+ )
+ for is_upstream, enabled in ((True, upstream), (False, downstream)):
+ if not enabled:
+ continue
+ for ti in tuple(retry.values()):
+ task = resolver.get_task(dag_id, run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+ contexts = resolver.producer_contexts(ti)
+ count = (
+ get_mapped_ti_count(task, run_id, session=session, producer_contexts=contexts)
+ if task.get_needs_expansion() and ti.region_index >= 0
+ else None
+ )
+ for relative in task.get_flat_relatives(upstream=is_upstream):
+ indexes = _get_relevant_map_indexes(
+ task=task,
+ run_id=run_id,
+ map_index=resolver.public_map_index(ti),
+ relative=relative,
+ ti_count=count,
+ session=session,
+ producer_contexts=contexts,
+ )
+ relative_loop = enclosing_loop(relative)
+ task_loop = enclosing_loop(task)
+ matches: Collection[TaskInstance]
+ if relative_loop is not None and (
+ task_loop is None or task_loop.group_id != relative_loop.group_id
+ ):
+ matches = [
+ other
+ for other in live.values()
+ if other.task_id == relative.task_id
+ and (
+ indexes is None
+ or (isinstance(indexes, int) and resolver.public_map_index(other) == indexes)
+ or (isinstance(indexes, range) and resolver.public_map_index(other) in indexes)
+ )
+ ]
+ else:
+ matches = resolver.resolve(
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=relative.task_id,
+ caller=ti,
+ map_indexes=indexes,
+ )
+ retry.update((other.id, other) for other in matches)
+ archived: set[UUID] = set()
+ if later_loop_iterations:
+ for ti in retry.values():
+ task = resolver.get_task(dag_id, run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+ group = enclosing_loop(task)
+ if group is None or ti.task_id != group.gate_task_id:
+ continue
+ position = loop_position(regions, ti.region_id, ti.region_index, group.node_id)
+ if position is None:
+ raise ValueError("Gate coordinates do not belong to its pinned loop")
+ for other in live.values():
+ other_position = loop_position(regions, other.region_id, other.region_index, group.node_id)
+ if (
+ other_position is not None
+ and other_position[0] == position[0]
+ and other_position[1] > position[1]
+ ):
+ archived.add(other.id)
+ return LoopClearScope(frozenset(retry.keys() - archived), frozenset(archived))
+
+
+def _regions_for_run(ti: TaskInstance, session: Session) -> dict[UUID, DynamicRegion]:
+ return {
+ region.id: region
+ for region in session.scalars(
+ select(DynamicRegion).where(DynamicRegion.dag_id == ti.dag_id, DynamicRegion.run_id == ti.run_id)
+ )
+ }
+
+
+def _physical_loop_coordinate(
+ region_id: UUID, index: int, node_id: str, regions: dict[UUID, DynamicRegion]
+) -> tuple[UUID, int] | None:
+ while region_id in regions:
+ region = regions[region_id]
+ if region.node_id == node_id:
+ return region_id, index
+ if region.parent_region_id is None or region.parent_region_index is None:
+ return None
+ region_id, index = region.parent_region_id, region.parent_region_index
+ return None
+
+
+def _forks_after(region_id: UUID, regions: dict[UUID, DynamicRegion]) -> Iterator[DynamicRegion]:
+ successors = {region.forked_from_region_id: region for region in regions.values()}
+ while region_id in successors:
+ successor = successors[region_id]
+ yield successor
+ region_id = successor.id
+
+
+def loop_execution_is_superseded(ti: TaskInstance, *, session: Session) -> bool:
+ """Identify a archiving loop coordinate for termination acknowledgement."""
+ return loop_coordinate_is_superseded(ti.region_id, ti.region_index, _regions_for_run(ti, session))
+
+
+def loop_coordinate_is_superseded(region_id: UUID, index: int, regions: dict[UUID, DynamicRegion]) -> bool:
+ """Test whether a coordinate or its parent lies beyond a durable fork cut."""
+ while region_id in regions:
+ if any(region.resumes_from_index <= index for region in _forks_after(region_id, regions)):
+ return True
+ region = regions[region_id]
+ if region.parent_region_id is None or region.parent_region_index is None:
+ break
+ region_id, index = region.parent_region_id, region.parent_region_index
+ return False
+
+
+def loop_gate_waits_for_archival(gate: TaskInstance, *, session: Session) -> bool:
+ """Hold a rerun gate until superseded executions of its later passes finish termination."""
+ if gate.operator != "LoopGateOperator" or gate.region_id == SENTINEL_REGION_ID:
+ return False
+ if not session.scalar(
+ select(
+ exists().where(
+ DynamicRegion.dag_id == gate.dag_id,
+ DynamicRegion.run_id == gate.run_id,
+ DynamicRegion.forked_from_region_id.is_not(None),
+ )
+ )
+ ):
+ return False
+ regions = _regions_for_run(gate, session)
+ gate_region = regions.get(gate.region_id)
+ if gate_region is None:
+ return False
+ node_id = gate_region.node_id
+ gate_position = loop_position(regions, gate.region_id, gate.region_index, node_id)
+ if gate_position is None:
+ raise ValueError("Gate coordinates do not belong to its pinned loop")
+ for region_id, region_index in session.execute(
+ select(TaskInstance.region_id, TaskInstance.region_index).where(
+ TaskInstance.dag_id == gate.dag_id,
+ TaskInstance.run_id == gate.run_id,
+ TaskInstance.working_set.is_(True),
+ TaskInstance.state == TaskInstanceState.RESTARTING,
+ )
+ ):
+ position = loop_position(regions, region_id, region_index, node_id)
+ if position is None or position[0] != gate_position[0] or position[1] <= gate_position[1]:
+ continue
+ coordinate = _physical_loop_coordinate(region_id, region_index, node_id, regions)
+ if coordinate is not None and any(
+ region.resumes_from_index <= coordinate[1] for region in _forks_after(coordinate[0], regions)
+ ):
+ return True
+ return False
+
+
+def loop_gate_has_later_pass(gate: TaskInstance, group: SerializedLoopTaskGroup, *, session: Session) -> bool:
+ """Find a live gate of a later pass in the same loop family, including passes in forked regions."""
+ regions = _regions_for_run(gate, session)
+ position = loop_position(regions, gate.region_id, gate.region_index, group.node_id)
+ if position is None:
+ raise ValueError("Gate coordinates do not belong to its pinned loop")
+ for region_id, region_index in session.execute(
+ select(TaskInstance.region_id, TaskInstance.region_index).where(
+ TaskInstance.working_set.is_(True),
+ TaskInstance.dag_id == gate.dag_id,
+ TaskInstance.run_id == gate.run_id,
+ TaskInstance.task_id == gate.task_id,
+ TaskInstance.region_index > gate.region_index,
+ )
+ ):
+ other = loop_position(regions, region_id, region_index, group.node_id)
+ if other is not None and other[0] == position[0]:
+ return True
+ return False
+
+
+def loop_successor_region(gate: TaskInstance, *, session: Session) -> UUID:
+ """Choose the region for a newly generated successor pass."""
+ region_id = gate.region_id
+ for fork in _forks_after(gate.region_id, _regions_for_run(gate, session)):
+ if fork.resumes_from_index <= gate.region_index + 1:
+ region_id = fork.id
+ return region_id
+
+
+def clear_loop_task_instances(
+ selected: Collection[TaskInstance],
+ *,
+ whole_expansion_ids: Collection[UUID] = (),
+ upstream: bool = False,
+ downstream: bool = True,
+ later_loop_iterations: bool = True,
+ session: Session,
+ dag_run_state: DagRunState | Literal[False] = DagRunState.QUEUED,
+ run_on_latest_version: bool = False,
+ prevent_running_task: bool | None = None,
+ whole_task_keys: Collection[tuple[str, str, str]] = (),
+) -> LoopClearScope:
+ """Clear selected loop tries and archive their replaced generated suffixes."""
+ from airflow.models.dagrun import DagRun
+
+ if not selected:
+ return LoopClearScope(frozenset(), frozenset())
+ run_keys = {(ti.dag_id, ti.run_id) for ti in selected}
+ if len(run_keys) != 1:
+ raise ValueError("Loop clear selection must belong to one DagRun")
+ dag_id, run_id = run_keys.pop()
+ session.scalar(select(DagRun).where(DagRun.dag_id == dag_id, DagRun.run_id == run_id).with_for_update())
+ scope = select_loop_clear_scope(
+ selected,
+ whole_expansion_ids=whole_expansion_ids,
+ upstream=upstream,
+ downstream=downstream,
+ later_loop_iterations=later_loop_iterations,
+ session=session,
+ )
+ apply_loop_clear_scope(
+ scope,
+ later_loop_iterations=later_loop_iterations,
+ session=session,
+ dag_run_state=dag_run_state,
+ run_on_latest_version=run_on_latest_version,
+ prevent_running_task=prevent_running_task,
+ whole_task_keys=whole_task_keys,
+ )
+ return scope
+
+
+def clear_task_instances_for_runs(
+ tis: Collection[TaskInstance],
+ *,
+ session: Session,
+ dag_run_state: DagRunState | Literal[False] = DagRunState.QUEUED,
+ run_on_latest_version: bool = False,
+ prevent_running_task: bool | None = None,
+ whole_task_keys: Collection[tuple[str, str, str]] = (),
+ later_loop_iterations: bool = True,
+) -> None:
+ """
+ Clear exactly the given executions, archiving the later passes of any loop gate among them.
+
+ Runs without a selected gate are retried in place as an ordinary clear.
+ """
+ from airflow.models.dagrun import DagRun
+
+ if not tis:
+ return
+ run_keys = sorted({(ti.dag_id, ti.run_id) for ti in tis})
+ session.scalars(
+ select(DagRun.id)
+ .where(tuple_(DagRun.dag_id, DagRun.run_id).in_(run_keys))
+ .order_by(DagRun.dag_id, DagRun.run_id)
+ .with_for_update()
+ ).all()
+ gate_runs = {(ti.dag_id, ti.run_id) for ti in tis if ti.operator == "LoopGateOperator"}
+ clear_task_instances(
+ [ti for ti in tis if (ti.dag_id, ti.run_id) not in gate_runs],
+ session=session,
+ dag_run_state=dag_run_state,
+ run_on_latest_version=run_on_latest_version,
+ prevent_running_task=prevent_running_task,
+ whole_task_keys=whole_task_keys,
+ )
+ for run_key in sorted(gate_runs):
+ clear_loop_task_instances(
+ [ti for ti in tis if (ti.dag_id, ti.run_id) == run_key],
+ downstream=False,
+ later_loop_iterations=later_loop_iterations,
+ session=session,
+ dag_run_state=dag_run_state,
+ run_on_latest_version=run_on_latest_version,
+ prevent_running_task=prevent_running_task,
+ whole_task_keys=whole_task_keys,
+ )
+
+
+def apply_loop_clear_scope(
+ scope: LoopClearScope,
+ *,
+ later_loop_iterations: bool = True,
+ session: Session,
+ dag_run_state: DagRunState | Literal[False] = DagRunState.QUEUED,
+ run_on_latest_version: bool = False,
+ prevent_running_task: bool | None = None,
+ whole_task_keys: Collection[tuple[str, str, str]] = (),
+) -> list[TaskInstance]:
+ """Apply a planned clear while the caller holds its DagRun lock."""
+ selected_ids = scope.retry_ids | scope.archive_ids
+ if not selected_ids:
+ return []
+ tis = list(
+ session.scalars(
+ select(TaskInstance)
+ .where(TaskInstance.id.in_(selected_ids), TaskInstance.working_set.is_(True))
+ .execution_options(populate_existing=True)
+ )
+ )
+ if {ti.id for ti in tis} != selected_ids:
+ raise ValueError("Clear selection contains executions that are no longer live")
+ run_keys = {(ti.dag_id, ti.run_id) for ti in tis}
+ if len(run_keys) != 1:
+ raise ValueError("Clear selection must belong to one DagRun")
+ dag_id, run_id = run_keys.pop()
+ if prevent_running_task and any(
+ ti.state in (TaskInstanceState.RUNNING, TaskInstanceState.RESTARTING) for ti in tis
+ ):
+ raise AirflowClearRunningTaskException("Cannot clear running task instances")
+ if later_loop_iterations:
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+ regions = _regions_for_run(tis[0], session)
+ cuts: dict[UUID, int] = {}
+ for ti in tis:
+ if ti.id not in scope.retry_ids or ti.operator != "LoopGateOperator":
+ continue
+ task = resolver.get_task(dag_id, run_id, ti.task_id, dag_version_id=ti.dag_version_id)
+ group = enclosing_loop(task)
+ if group is None or group.gate_task_id != ti.task_id:
+ raise ValueError("Gate coordinates require a pinned loop definition")
+ family = loop_position(regions, ti.region_id, ti.region_index, group.node_id)
+ if family is None:
+ raise ValueError("Gate coordinates do not belong to its pinned loop")
+ cuts[family[0]] = min(cuts.get(family[0], ti.region_index + 1), ti.region_index + 1)
+ for family_id, index in cuts.items():
+ leaf = regions[family_id]
+ for successor in _forks_after(family_id, regions):
+ leaf = successor
+ session.add(
+ DynamicRegion(
+ dag_id=dag_id,
+ run_id=run_id,
+ node_id=leaf.node_id,
+ parent_region_id=leaf.parent_region_id,
+ parent_region_index=leaf.parent_region_index,
+ forked_from_region_id=leaf.id,
+ resumes_from_index=index,
+ )
+ )
+ session.flush()
+ return clear_task_instances(
+ tis,
+ session=session,
+ superseded_ti_ids=scope.archive_ids,
+ dag_run_state=dag_run_state,
+ run_on_latest_version=run_on_latest_version,
+ prevent_running_task=prevent_running_task,
+ whole_task_keys=whole_task_keys,
+ )
diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py
index 4891e103a32..06510a7dbe6 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -400,6 +400,9 @@ def clear_task_instances(
dag_run_state: DagRunState | Literal[False] = DagRunState.QUEUED,
run_on_latest_version: bool = False,
prevent_running_task: bool | None = None,
+ *,
+ whole_task_keys: Collection[tuple[str, str, str]] = (),
+ superseded_ti_ids: Collection[UUID] = (),
) -> list[TaskInstance]:
"""
Clear a set of task instances, but make sure the running ones get killed.
@@ -417,11 +420,51 @@ def clear_task_instances(
:param run_on_latest_version: whether to run on latest serialized DAG and Bundle version.
A run with no version of its own uses the latest either way, since there is nothing
else for it to run on; a task instance with no version joins its run's.
+ :param whole_task_keys: ``(dag_id, run_id, task_id)`` of tasks cleared as a whole. A pre-region
+ mapped expansion of such a task is archived and regenerated in a new region.
+ :param superseded_ti_ids: executions whose generated work is replaced rather than retried; they
+ are archived in place instead of receiving a successor.
:meta private:
"""
from airflow.exceptions import AirflowClearRunningTaskException
from airflow.models.dagbag import DBDagBag
+ from airflow.models.dagrun import DagRun
+
+ if tis:
+ run_keys = {(ti.dag_id, ti.run_id) for ti in tis}
+ session.scalars(
+ select(DagRun.id)
+ .where(tuple_(DagRun.dag_id, DagRun.run_id).in_(run_keys))
+ .order_by(DagRun.dag_id, DagRun.run_id)
+ .with_for_update()
+ ).all()
+
+ legacy_keys = {
+ (ti.dag_id, ti.run_id, ti.task_id)
+ for ti in tis
+ if ti.region_id == SENTINEL_REGION_ID
+ and ti.region_index >= 0
+ and (ti.dag_id, ti.run_id, ti.task_id) in whole_task_keys
+ }
+ if legacy_keys:
+ existing_keys = {
+ (region.dag_id, region.run_id, region.node_id)
+ for region in session.scalars(
+ select(DynamicRegion).where(
+ tuple_(DynamicRegion.dag_id, DynamicRegion.run_id, DynamicRegion.node_id).in_(
+ legacy_keys
+ ),
+ DynamicRegion.parent_region_id.is_(None),
+ DynamicRegion.forked_from_region_id.is_(None),
+ )
+ )
+ }
+ session.add_all(
+ DynamicRegion(dag_id=dag_id, run_id=run_id, node_id=task_id)
+ for dag_id, run_id, task_id in legacy_keys - existing_keys
+ )
+ session.flush()
scheduler_dagbag = DBDagBag(load_op_links=False)
cleared = []
@@ -440,6 +483,10 @@ def clear_task_instances(
# If a task is cleared when running and the prevent_running_task is false,
# set its state to RESTARTING so that
# the task is terminated and becomes eligible for retry.
+ elif ti.id in superseded_ti_ids or (
+ ti.region_id == SENTINEL_REGION_ID and (ti.dag_id, ti.run_id, ti.task_id) in legacy_keys
+ ):
+ ti.archive(reason="superseded", session=session)
else:
if ti.state in (None, TaskInstanceState.UP_FOR_RETRY):
# The pending attempt hasn't run, so base its retry budget on the preceding attempt.
@@ -568,7 +615,8 @@ def clear_task_instances(
if dr.created_dag_version_id:
_pin_versionless_tis_to_run_version(dr, dr.created_dag_version_id, session)
for ti in tis:
- ti.context_carrier = new_task_run_carrier(ti.dag_run.context_carrier)
+ if ti.working_set is True:
+ ti.context_carrier = new_task_run_carrier(ti.dag_run.context_carrier)
session.flush()
return cleared
@@ -1212,14 +1260,45 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
# is the task still in the retry waiting period?
return self.state == TaskInstanceState.UP_FOR_RETRY and not self.ready_for_retry()
+ def is_replaced_generated_work(self, *, session: Session) -> bool:
+ """Whether a pre-region mapped expansion was cleared as a whole and now has a replacement region."""
+ if self.region_id != SENTINEL_REGION_ID or self.region_index < 0:
+ return False
+ return (
+ session.scalar(
+ select(DynamicRegion.id)
+ .where(
+ DynamicRegion.dag_id == self.dag_id,
+ DynamicRegion.run_id == self.run_id,
+ DynamicRegion.node_id == self.task_id,
+ DynamicRegion.parent_region_id.is_(None),
+ DynamicRegion.forked_from_region_id.is_(None),
+ )
+ .limit(1)
+ )
+ is not None
+ )
+
def archive(self, *, reason: str, session: Session) -> None:
"""Remove this attempt from the working set while retaining its UUID and children."""
- current = session.scalar(
- select(TaskInstance.working_set).where(TaskInstance.id == self.id).with_for_update()
- )
- if current is not True:
+ current = session.execute(
+ select(
+ TaskInstance.working_set,
+ TaskInstance.state,
+ TaskInstance.start_date,
+ TaskInstance.end_date,
+ )
+ .where(TaskInstance.id == self.id)
+ .with_for_update()
+ ).one_or_none()
+ if current is None or current.working_set is not True:
raise ValueError("An archived task instance cannot be archived again")
- if self.state not in State.finished:
+ state: Any = inspect(self)
+ for name in ("state", "start_date", "end_date"):
+ if not state.attrs[name].history.has_changes():
+ setattr(self, name, getattr(current, name))
+ never_started = reason == "superseded" and self.start_date is None
+ if self.state not in State.finished and not never_started:
self.state = TaskInstanceState.FAILED
if self.end_date is None:
self.end_date = timezone.utcnow()
@@ -1308,10 +1387,30 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
successor.task_instance_note = carried
return successor
- def complete_restart(self, *, session: Session) -> TaskInstance:
- """Release a cleared attempt after termination; the caller must hold its row lock."""
+ def complete_restart(
+ self, *, session: Session, terminal_outcome: TaskInstanceState | None = None
+ ) -> TaskInstance:
+ """
+ Release a cleared attempt after termination; the caller must hold its row lock.
+
+ A task instance whose generated work was replaced is archived in place instead of receiving a
+ successor, ending in ``terminal_outcome`` when one is given.
+ """
+ from airflow.models.loop_clear import loop_execution_is_superseded
+
if self.state != TaskInstanceState.RESTARTING or self.working_set is not True:
raise ValueError("Only a current restarting task instance can complete a restart")
+ if terminal_outcome is not None and terminal_outcome not in State.finished:
+ raise ValueError("An archival outcome must be terminal")
+ if self.is_replaced_generated_work(session=session) or (
+ self.region_id != SENTINEL_REGION_ID and loop_execution_is_superseded(self, session=session)
+ ):
+ if terminal_outcome is not None:
+ self.set_state(terminal_outcome, session=session)
+ self.archive(reason="superseded", session=session)
+ return self
+ if terminal_outcome is not None:
+ raise ValueError("A terminal archival outcome requires a superseded execution")
successor = self.prepare_db_for_next_try(session)
if self.task is not None:
successor.max_tries = self.try_number + self.task.retries
diff --git a/airflow-core/src/airflow/serialization/definitions/baseoperator.py b/airflow-core/src/airflow/serialization/definitions/baseoperator.py
index 8c28b17c377..3fadf781d61 100644
--- a/airflow-core/src/airflow/serialization/definitions/baseoperator.py
+++ b/airflow-core/src/airflow/serialization/definitions/baseoperator.py
@@ -28,6 +28,7 @@ from airflow.serialization.definitions.node import DAGNode
from airflow.serialization.definitions.param import SerializedParamsDict
from airflow.serialization.enums import DagAttributeTypes
from airflow.task.priority_strategy import PriorityWeightStrategy, validate_and_load_priority_weight_strategy
+from airflow.ti_deps.deps.loop_archival_dep import LoopArchivalDep
from airflow.ti_deps.deps.mapped_task_upstream_dep import MappedTaskUpstreamDep
from airflow.ti_deps.deps.not_in_retry_period_dep import NotInRetryPeriodDep
from airflow.ti_deps.deps.not_previously_skipped_dep import NotPreviouslySkippedDep
@@ -54,6 +55,7 @@ DEFAULT_OPERATOR_DEPS: frozenset[BaseTIDep] = frozenset(
TriggerRuleDep(),
NotPreviouslySkippedDep(),
MappedTaskUpstreamDep(),
+ LoopArchivalDep(),
)
)
diff --git a/airflow-core/src/airflow/serialization/definitions/dag.py b/airflow-core/src/airflow/serialization/definitions/dag.py
index 041cc88a358..580a6b1de61 100644
--- a/airflow-core/src/airflow/serialization/definitions/dag.py
+++ b/airflow-core/src/airflow/serialization/definitions/dag.py
@@ -1225,10 +1225,10 @@ class SerializedDAG:
tuples that should not be cleared
:param exclude_run_ids: A set of ``run_id`` or (``run_id``)
"""
+ from airflow.models.loop_clear import clear_task_instances_for_runs
from airflow.models.taskinstance import (
_get_new_task_ids,
_update_dagrun_to_latest_version,
- clear_task_instances,
)
if only_new:
@@ -1290,11 +1290,23 @@ class SerializedDAG:
if count == 0:
return 0
- clear_task_instances(
- list(tis),
- session,
+ clear_task_instances_for_runs(
+ tis,
+ session=session,
dag_run_state=dag_run_state,
run_on_latest_version=run_on_latest_version,
+ later_loop_iterations=not (only_failed or only_running),
+ whole_task_keys={
+ (ti.dag_id, ti.run_id, ti.task_id)
+ for ti in tis
+ if not only_failed
+ and not only_running
+ and (task_ids is None or ti.task_id in task_ids)
+ and not any(
+ isinstance(excluded, tuple) and excluded[0] == ti.task_id
+ for excluded in exclude_task_ids or ()
+ )
+ },
)
session.flush()
diff --git a/airflow-core/src/airflow/ti_deps/deps/loop_archival_dep.py b/airflow-core/src/airflow/ti_deps/deps/loop_archival_dep.py
new file mode 100644
index 00000000000..d688c8b8c3d
--- /dev/null
+++ b/airflow-core/src/airflow/ti_deps/deps/loop_archival_dep.py
@@ -0,0 +1,32 @@
+# 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 airflow.ti_deps.deps.base_ti_dep import BaseTIDep
+
+
+class LoopArchivalDep(BaseTIDep):
+ """A rerun gate waits for the generated work it replaces to terminate."""
+
+ NAME = "Loop archival"
+ IS_TASK_DEP = True
+
+ def _get_dep_statuses(self, ti, dep_context, session):
+ from airflow.models.loop_clear import loop_gate_waits_for_archival
+
+ if loop_gate_waits_for_archival(ti, session=session):
+ yield self._failing_status(reason="Superseded loop executions are still terminating")
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
index a81337ba5b3..af1c43e4301 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -2898,6 +2898,8 @@ export const useDagServiceBulkDags = <TData = Common.DagServiceBulkDagsMutationR
* @param data.requestBody
* @param data.mapIndex
* @param data.updateMask
+* @param data.regionId
+* @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -2905,6 +2907,8 @@ export const useTaskInstanceServicePatchTaskInstance = <TData = Common.TaskInsta
dagId: string;
dagRunId: string;
mapIndex?: number;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: string[];
@@ -2912,10 +2916,12 @@ export const useTaskInstanceServicePatchTaskInstance = <TData = Common.TaskInsta
dagId: string;
dagRunId: string;
mapIndex?: number;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: string[];
-}, TContext>({ mutationFn: ({ dagId, dagRunId, mapIndex, requestBody, taskId, updateMask }) => TaskInstanceService.patchTaskInstance({ dagId, dagRunId, mapIndex, requestBody, taskId, updateMask }) as unknown as Promise<TData>, ...options });
+}, TContext>({ mutationFn: ({ dagId, dagRunId, mapIndex, regionId, regionIndex, requestBody, taskId, updateMask }) => TaskInstanceService.patchTaskInstance({ dagId, dagRunId, mapIndex, regionId, regionIndex, requestBody, taskId, updateMask }) as unknown as Promise<TData>, ...options });
/**
* Patch Task Instance
* Update a task instance.
@@ -2926,6 +2932,8 @@ export const useTaskInstanceServicePatchTaskInstance = <TData = Common.TaskInsta
* @param data.mapIndex
* @param data.requestBody
* @param data.updateMask
+* @param data.regionId
+* @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -2933,6 +2941,8 @@ export const useTaskInstanceServicePatchTaskInstanceByMapIndex = <TData = Common
dagId: string;
dagRunId: string;
mapIndex: number;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: string[];
@@ -2940,10 +2950,12 @@ export const useTaskInstanceServicePatchTaskInstanceByMapIndex = <TData = Common
dagId: string;
dagRunId: string;
mapIndex: number;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: string[];
-}, TContext>({ mutationFn: ({ dagId, dagRunId, mapIndex, requestBody, taskId, updateMask }) => TaskInstanceService.patchTaskInstanceByMapIndex({ dagId, dagRunId, mapIndex, requestBody, taskId, updateMask }) as unknown as Promise<TData>, ...options });
+}, TContext>({ mutationFn: ({ dagId, dagRunId, mapIndex, regionId, regionIndex, requestBody, taskId, updateMask }) => TaskInstanceService.patchTaskInstanceByMapIndex({ dagId, dagRunId, mapIndex, regionId, regionIndex, requestBody, taskId, updateMask }) as unknown as Promise<TData>, ...options });
/**
* Bulk Task Instances
* Bulk update, and delete task instances.
@@ -2972,6 +2984,8 @@ export const useTaskInstanceServiceBulkTaskInstances = <TData = Common.TaskInsta
* @param data.groupId
* @param data.requestBody
* @param data.updateMask
+* @param data.regionId
+* @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -2979,15 +2993,19 @@ export const useTaskInstanceServicePatchTaskGroupInstances = <TData = Common.Tas
dagId: string;
dagRunId: string;
groupId: string;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
updateMask?: string[];
}, TContext>, "mutationFn">) => useMutation<TData, TError, {
dagId: string;
dagRunId: string;
groupId: string;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
updateMask?: string[];
-}, TContext>({ mutationFn: ({ dagId, dagRunId, groupId, requestBody, updateMask }) => TaskInstanceService.patchTaskGroupInstances({ dagId, dagRunId, groupId, requestBody, updateMask }) as unknown as Promise<TData>, ...options });
+}, TContext>({ mutationFn: ({ dagId, dagRunId, groupId, regionId, regionIndex, requestBody, updateMask }) => TaskInstanceService.patchTaskGroupInstances({ dagId, dagRunId, groupId, regionId, regionIndex, requestBody, updateMask }) as unknown as Promise<TData>, ...options });
/**
* Patch Task Group Instances Dry Run
* Dry-run of updating the state of all task instances in a task group.
@@ -2996,6 +3014,8 @@ export const useTaskInstanceServicePatchTaskGroupInstances = <TData = Common.Tas
* @param data.dagRunId
* @param data.groupId
* @param data.requestBody
+* @param data.regionId
+* @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -3003,13 +3023,17 @@ export const useTaskInstanceServicePatchTaskGroupInstancesDryRun = <TData = Comm
dagId: string;
dagRunId: string;
groupId: string;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
}, TContext>, "mutationFn">) => useMutation<TData, TError, {
dagId: string;
dagRunId: string;
groupId: string;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
-}, TContext>({ mutationFn: ({ dagId, dagRunId, groupId, requestBody }) => TaskInstanceService.patchTaskGroupInstancesDryRun({ dagId, dagRunId, groupId, requestBody }) as unknown as Promise<TData>, ...options });
+}, TContext>({ mutationFn: ({ dagId, dagRunId, groupId, regionId, regionIndex, requestBody }) => TaskInstanceService.patchTaskGroupInstancesDryRun({ dagId, dagRunId, groupId, regionId, regionIndex, requestBody }) as unknown as Promise<TData>, ...options });
/**
* Patch Task Instance Dry Run
* Update a task instance dry_run mode.
@@ -3020,6 +3044,8 @@ export const useTaskInstanceServicePatchTaskGroupInstancesDryRun = <TData = Comm
* @param data.mapIndex
* @param data.requestBody
* @param data.updateMask
+* @param data.regionId
+* @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -3027,6 +3053,8 @@ export const useTaskInstanceServicePatchTaskInstanceDryRunByMapIndex = <TData =
dagId: string;
dagRunId: string;
mapIndex: number;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: string[];
@@ -3034,10 +3062,12 @@ export const useTaskInstanceServicePatchTaskInstanceDryRunByMapIndex = <TData =
dagId: string;
dagRunId: string;
mapIndex: number;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: string[];
-}, TContext>({ mutationFn: ({ dagId, dagRunId, mapIndex, requestBody, taskId, updateMask }) => TaskInstanceService.patchTaskInstanceDryRunByMapIndex({ dagId, dagRunId, mapIndex, requestBody, taskId, updateMask }) as unknown as Promise<TData>, ...options });
+}, TContext>({ mutationFn: ({ dagId, dagRunId, mapIndex, regionId, regionIndex, requestBody, taskId, updateMask }) => TaskInstanceService.patchTaskInstanceDryRunByMapIndex({ dagId, dagRunId, mapIndex, regionId, regionIndex, requestBody, taskId, updateMask }) as unknown as Promise<TData>, ...options });
/**
* Patch Task Instance Dry Run
* Update a task instance dry_run mode.
@@ -3048,6 +3078,8 @@ export const useTaskInstanceServicePatchTaskInstanceDryRunByMapIndex = <TData =
* @param data.requestBody
* @param data.mapIndex
* @param data.updateMask
+* @param data.regionId
+* @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -3055,6 +3087,8 @@ export const useTaskInstanceServicePatchTaskInstanceDryRun = <TData = Common.Tas
dagId: string;
dagRunId: string;
mapIndex?: number;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: string[];
@@ -3062,10 +3096,12 @@ export const useTaskInstanceServicePatchTaskInstanceDryRun = <TData = Common.Tas
dagId: string;
dagRunId: string;
mapIndex?: number;
+ regionId?: string;
+ regionIndex?: number;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: string[];
-}, TContext>({ mutationFn: ({ dagId, dagRunId, mapIndex, requestBody, taskId, updateMask }) => TaskInstanceService.patchTaskInstanceDryRun({ dagId, dagRunId, mapIndex, requestBody, taskId, updateMask }) as unknown as Promise<TData>, ...options });
+}, TContext>({ mutationFn: ({ dagId, dagRunId, mapIndex, regionId, regionIndex, requestBody, taskId, updateMask }) => TaskInstanceService.patchTaskInstanceDryRun({ dagId, dagRunId, mapIndex, regionId, regionIndex, requestBody, taskId, updateMask }) as unknown as Promise<TData>, ...options });
/**
* Update Hitl Detail
* Update a Human-in-the-loop detail.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
index 2c577c6fed2..63947f8f223 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -1912,6 +1912,29 @@ export const $BulkTaskInstanceBody = {
}
]
},
+ region_id: {
+ anyOf: [
+ {
+ type: 'string',
+ format: 'uuid'
+ },
+ {
+ type: 'null'
+ }
+ ],
+ title: 'Region Id'
+ },
+ region_index: {
+ anyOf: [
+ {
+ type: 'integer'
+ },
+ {
+ type: 'null'
+ }
+ ],
+ title: 'Region Index'
+ },
note: {
anyOf: [
{
@@ -2450,6 +2473,38 @@ export const $ClearTaskInstancesBody = {
title: 'Task Group Id',
description: "Clear every task in this task group. Mutually exclusive with `task_ids`. The group's tasks are resolved on the server from the dag structure, so all of them are targeted regardless of how many there are."
},
+ task_instance_ids: {
+ anyOf: [
+ {
+ items: {
+ type: 'string',
+ format: 'uuid'
+ },
+ type: 'array',
+ minItems: 1
+ },
+ {
+ type: 'null'
+ }
+ ],
+ title: 'Task Instance Ids',
+ description: 'Exact current executions in one DAG run. Mutually exclusive with task_ids and task_group_id.'
+ },
+ whole_expansion_ids: {
+ items: {
+ type: 'string',
+ format: 'uuid'
+ },
+ type: 'array',
+ title: 'Whole Expansion Ids',
+ description: 'Selected execution IDs whose entire mapped expansion should be cleared.'
+ },
+ include_later_loop_iterations: {
+ type: 'boolean',
+ title: 'Include Later Loop Iterations',
+ description: 'Clear later iterations when the selection includes a loop gate.',
+ default: true
+ },
dag_run_id: {
anyOf: [
{
@@ -6785,6 +6840,29 @@ export const $PatchTaskInstanceBody = {
}
]
},
+ region_id: {
+ anyOf: [
+ {
+ type: 'string',
+ format: 'uuid'
+ },
+ {
+ type: 'null'
+ }
+ ],
+ title: 'Region Id'
+ },
+ region_index: {
+ anyOf: [
+ {
+ type: 'integer'
+ },
+ {
+ type: 'null'
+ }
+ ],
+ title: 'Region Index'
+ },
note: {
anyOf: [
{
@@ -8007,6 +8085,15 @@ export const $TaskInstanceResponse = {
type: 'integer',
title: 'Map Index'
},
+ region_id: {
+ type: 'string',
+ format: 'uuid',
+ title: 'Region Id'
+ },
+ region_index: {
+ type: 'integer',
+ title: 'Region Index'
+ },
logical_date: {
anyOf: [
{
@@ -8291,7 +8378,7 @@ export const $TaskInstanceResponse = {
}
},
type: 'object',
- required: ['id', 'task_id', 'dag_id', 'dag_run_id', 'map_index', 'logical_date', 'run_after', 'start_date', 'end_date', 'duration', 'state', 'try_number', 'max_tries', 'task_display_name', 'dag_display_name', 'hostname', 'unixname', 'pool', 'pool_slots', 'queue', 'priority_weight', 'operator', 'operator_name', 'queued_when', 'scheduled_when', 'pid', 'executor', 'executor_config', 'note', 'rendered_map_index', 'trigger', 'triggerer_job', 'dag_version'],
+ required: ['id', 'task_id', 'dag_id', 'dag_run_id', 'map_index', 'region_id', 'region_index', 'logical_date', 'run_after', 'start_date', 'end_date', 'duration', 'state', 'try_number', 'max_tries', 'task_display_name', 'dag_display_name', 'hostname', 'unixname', 'pool', 'pool_slots', 'queue', 'priority_weight', 'operator', 'operator_name', 'queued_when', 'scheduled_when', 'pid', 'executor', 'executor_config', 'note', 'rendered_map_index', 'trigger', 'triggerer_job', 'dag_version'],
title: 'TaskInstanceResponse',
description: 'TaskInstance serializer for responses.'
} as const;
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
index 6c8db6ea6ed..9fdae3288dc 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -2581,6 +2581,8 @@ export class TaskInstanceService {
* @param data.requestBody
* @param data.mapIndex
* @param data.updateMask
+ * @param data.regionId
+ * @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -2595,7 +2597,9 @@ export class TaskInstanceService {
},
query: {
map_index: data.mapIndex,
- update_mask: data.updateMask
+ update_mask: data.updateMask,
+ region_id: data.regionId,
+ region_index: data.regionIndex
},
body: data.requestBody,
mediaType: 'application/json',
@@ -2921,6 +2925,8 @@ export class TaskInstanceService {
* @param data.mapIndex
* @param data.requestBody
* @param data.updateMask
+ * @param data.regionId
+ * @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -2935,7 +2941,9 @@ export class TaskInstanceService {
map_index: data.mapIndex
},
query: {
- update_mask: data.updateMask
+ update_mask: data.updateMask,
+ region_id: data.regionId,
+ region_index: data.regionIndex
},
body: data.requestBody,
mediaType: 'application/json',
@@ -3258,6 +3266,8 @@ export class TaskInstanceService {
* @param data.groupId
* @param data.requestBody
* @param data.updateMask
+ * @param data.regionId
+ * @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -3271,7 +3281,9 @@ export class TaskInstanceService {
group_id: data.groupId
},
query: {
- update_mask: data.updateMask
+ update_mask: data.updateMask,
+ region_id: data.regionId,
+ region_index: data.regionIndex
},
body: data.requestBody,
mediaType: 'application/json',
@@ -3294,6 +3306,8 @@ export class TaskInstanceService {
* @param data.dagRunId
* @param data.groupId
* @param data.requestBody
+ * @param data.regionId
+ * @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -3306,6 +3320,10 @@ export class TaskInstanceService {
dag_run_id: data.dagRunId,
group_id: data.groupId
},
+ query: {
+ region_id: data.regionId,
+ region_index: data.regionIndex
+ },
body: data.requestBody,
mediaType: 'application/json',
errors: {
@@ -3328,6 +3346,8 @@ export class TaskInstanceService {
* @param data.mapIndex
* @param data.requestBody
* @param data.updateMask
+ * @param data.regionId
+ * @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -3342,7 +3362,9 @@ export class TaskInstanceService {
map_index: data.mapIndex
},
query: {
- update_mask: data.updateMask
+ update_mask: data.updateMask,
+ region_id: data.regionId,
+ region_index: data.regionIndex
},
body: data.requestBody,
mediaType: 'application/json',
@@ -3366,6 +3388,8 @@ export class TaskInstanceService {
* @param data.requestBody
* @param data.mapIndex
* @param data.updateMask
+ * @param data.regionId
+ * @param data.regionIndex
* @returns TaskInstanceCollectionResponse Successful Response
* @throws ApiError
*/
@@ -3380,7 +3404,9 @@ export class TaskInstanceService {
},
query: {
map_index: data.mapIndex,
- update_mask: data.updateMask
+ update_mask: data.updateMask,
+ region_id: data.regionId,
+ region_index: data.regionIndex
},
body: data.requestBody,
mediaType: 'application/json',
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
index d0993f8d3f8..649895f93ae 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -593,6 +593,8 @@ export type BulkResponse = {
*/
export type BulkTaskInstanceBody = {
new_state?: TaskInstanceState | null;
+ region_id?: string | null;
+ region_index?: number | null;
note?: string | null;
include_upstream?: boolean;
include_downstream?: boolean;
@@ -768,6 +770,18 @@ export type ClearTaskInstancesBody = {
* Clear every task in this task group. Mutually exclusive with `task_ids`. The group's tasks are resolved on the server from the dag structure, so all of them are targeted regardless of how many there are.
*/
task_group_id?: string | null;
+ /**
+ * Exact current executions in one DAG run. Mutually exclusive with task_ids and task_group_id.
+ */
+ task_instance_ids?: Array<(string)> | null;
+ /**
+ * Selected execution IDs whose entire mapped expansion should be cleared.
+ */
+ whole_expansion_ids?: Array<(string)>;
+ /**
+ * Clear later iterations when the selection includes a loop gate.
+ */
+ include_later_loop_iterations?: boolean;
dag_run_id?: string | null;
include_upstream?: boolean;
include_downstream?: boolean;
@@ -1849,6 +1863,8 @@ export type NewTaskResponse = {
*/
export type PatchTaskInstanceBody = {
new_state?: TaskInstanceState | null;
+ region_id?: string | null;
+ region_index?: number | null;
note?: string | null;
include_upstream?: boolean;
include_downstream?: boolean;
@@ -2175,6 +2191,8 @@ export type TaskInstanceResponse = {
dag_id: string;
dag_run_id: string;
map_index: number;
+ region_id: string;
+ region_index: number;
logical_date: string | null;
run_after: string;
start_date: string | null;
@@ -4228,6 +4246,8 @@ export type PatchTaskInstanceData = {
dagId: string;
dagRunId: string;
mapIndex?: number | null;
+ regionId?: string | null;
+ regionIndex?: number | null;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: Array<(string)> | null;
@@ -4371,6 +4391,8 @@ export type PatchTaskInstanceByMapIndexData = {
dagId: string;
dagRunId: string;
mapIndex: number | null;
+ regionId?: string | null;
+ regionIndex?: number | null;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: Array<(string)> | null;
@@ -4536,6 +4558,8 @@ export type PatchTaskGroupInstancesData = {
dagId: string;
dagRunId: string;
groupId: string;
+ regionId?: string | null;
+ regionIndex?: number | null;
requestBody: PatchTaskInstanceBody;
updateMask?: Array<(string)> | null;
};
@@ -4546,6 +4570,8 @@ export type PatchTaskGroupInstancesDryRunData = {
dagId: string;
dagRunId: string;
groupId: string;
+ regionId?: string | null;
+ regionIndex?: number | null;
requestBody: PatchTaskInstanceBody;
};
@@ -4555,6 +4581,8 @@ export type PatchTaskInstanceDryRunByMapIndexData = {
dagId: string;
dagRunId: string;
mapIndex: number | null;
+ regionId?: string | null;
+ regionIndex?: number | null;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: Array<(string)> | null;
@@ -4566,6 +4594,8 @@ export type PatchTaskInstanceDryRunData = {
dagId: string;
dagRunId: string;
mapIndex?: number | null;
+ regionId?: string | null;
+ regionIndex?: number | null;
requestBody: PatchTaskInstanceBody;
taskId: string;
updateMask?: Array<(string)> | null;
diff --git a/airflow-core/src/airflow/ui/src/components/MarkAs/TaskInstance/MarkTaskInstanceAsDialog.test.tsx b/airflow-core/src/airflow/ui/src/components/MarkAs/TaskInstance/MarkTaskInstanceAsDialog.test.tsx
index 42cfc97c9e3..dea4e877ad0 100644
--- a/airflow-core/src/airflow/ui/src/components/MarkAs/TaskInstance/MarkTaskInstanceAsDialog.test.tsx
+++ b/airflow-core/src/airflow/ui/src/components/MarkAs/TaskInstance/MarkTaskInstanceAsDialog.test.tsx
@@ -66,6 +66,8 @@ const taskInstance: TaskInstanceResponse = {
priority_weight: null,
queue: null,
queued_when: null,
+ region_id: "00000000-0000-0000-0000-000000000000",
+ region_index: 0,
rendered_fields: undefined,
rendered_map_index: null,
run_after: "2025-01-01T00:00:00Z",
diff --git a/airflow-core/src/airflow/ui/src/components/TaskInstanceTooltip.test.tsx b/airflow-core/src/airflow/ui/src/components/TaskInstanceTooltip.test.tsx
index 3a5a5039d3e..5b725ab34c6 100644
--- a/airflow-core/src/airflow/ui/src/components/TaskInstanceTooltip.test.tsx
+++ b/airflow-core/src/airflow/ui/src/components/TaskInstanceTooltip.test.tsx
@@ -97,6 +97,8 @@ describe("TaskInstanceTooltip", () => {
priority_weight: null,
queue: null,
queued_when: null,
+ region_id: "00000000-0000-0000-0000-000000000000",
+ region_index: 0,
rendered_fields: undefined,
rendered_map_index: null,
run_after: twoHoursAgo,
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
index 68af746f035..a1efdfb95c5 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py
@@ -2255,6 +2255,40 @@ class TestClearDagRun:
)
assert logs == 0
+ @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
+ def test_clear_dag_run_dry_run_reports_loop_coordinates(self, test_client, dag_maker, session):
+ @task_group
+ def body():
+ EmptyOperator(task_id="work")
+
+ with dag_maker("dry_run_loop", serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ run = dag_maker.create_dagrun(run_id="dry_run_loop_run")
+ region = session.scalar(select(DynamicRegion).where(DynamicRegion.node_id == loop.group_id))
+ later = TaskInstance(
+ task=dag_maker.dag.get_task("body.work"),
+ run_id=run.run_id,
+ dag_version_id=run.created_dag_version_id,
+ region_id=region.id,
+ region_index=1,
+ )
+ later.state = State.SUCCESS
+ session.add(later)
+ session.commit()
+
+ response = test_client.post(
+ f"/dags/{run.dag_id}/dagRuns/{run.run_id}/clear", json={"dry_run": True, "only_failed": False}
+ )
+
+ assert response.status_code == 200, response.text
+ work = [row for row in response.json()["task_instances"] if row["task_id"] == "body.work"]
+ assert {row["region_index"] for row in work} == {0, 1}
+ assert {row["map_index"] for row in work} == {-1}
+ assert all(
+ row["loop_iterations"] == [{"loop_id": loop.group_id, "iteration": row["region_index"]}]
+ for row in work
+ )
+
@pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
def test_clear_dag_run_dry_run_response_has_full_task_instance_fields(self, test_client):
"""Regression test: dry-run response must include all TaskInstanceResponse fields.
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
index be5f88268fc..63cfe3c56ca 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
@@ -28,7 +28,8 @@ from unittest import mock
import pendulum
import pytest
from sqlalchemy import delete, func, select, update
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import Session, joinedload
+from sqlalchemy.sql.selectable import Select
from airflow._shared.secrets_masker import mask_secret
from airflow._shared.state import TaskScope
@@ -41,12 +42,16 @@ from airflow.jobs.triggerer_job_runner import TriggererJobRunner
from airflow.models import DagModel, DagRun, Log, TaskInstance
from airflow.models.dag_version import DagVersion
from airflow.models.dagbundle import DagBundleModel
+from airflow.models.dynamic_region import LOOP_DECISION_KEY, DynamicRegion
+from airflow.models.loop_clear import clear_loop_task_instances
from airflow.models.renderedtifields import RenderedTaskInstanceFields as RTIF
from airflow.models.task_state_store import TaskStateStoreModel
-from airflow.models.taskinstance import uuid7
+from airflow.models.taskinstance import clear_task_instances, uuid7
from airflow.models.team import Team
from airflow.models.trigger import Trigger
-from airflow.sdk import BaseOperator, TaskGroup, task
+from airflow.models.xcom import XComModelV2
+from airflow.sdk import BaseOperator, TaskGroup, task, task_group
+from airflow.sdk.definitions._internal.loop import create_loop
from airflow.state.metastore import MetastoreBackend
from airflow.utils.platform import getuser
from airflow.utils.state import DagRunState, State, TaskInstanceState
@@ -65,6 +70,7 @@ from tests_common.test_utils.mapping import expand_mapped_task_instances, push_m
from tests_common.test_utils.mock_operators import MockOperator
from tests_common.test_utils.taskinstance import create_task_instance
from tests_common.test_utils.team import attach_dag_to_team
+from unit.listeners.class_listener import ClassBasedListener
pytestmark = pytest.mark.db_test
@@ -174,7 +180,7 @@ class TestTaskInstanceEndpoint:
ti.try_number = 1
session.merge(ti)
session.flush()
- dag.clear(session=session)
+ clear_task_instances(tis, session=session)
successors = []
for ti in tis:
current = session.scalar(
@@ -261,6 +267,8 @@ class TestGetTaskInstance(TestTaskInstanceEndpoint):
"hostname": "",
"id": response_data["id"],
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"max_tries": 0,
"note": "placeholder-note",
"operator": None,
@@ -374,6 +382,8 @@ class TestGetTaskInstance(TestTaskInstanceEndpoint):
"dag_run_id": run_id,
"dag_display_name": "dag_with_multiple_versions",
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"logical_date": mock.ANY,
"start_date": None,
"end_date": mock.ANY,
@@ -462,6 +472,8 @@ class TestGetTaskInstance(TestTaskInstanceEndpoint):
"hostname": "",
"id": response_data["id"],
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"max_tries": 0,
"note": "placeholder-note",
"operator": "PythonOperator",
@@ -530,6 +542,8 @@ class TestGetTaskInstance(TestTaskInstanceEndpoint):
"hostname": "",
"id": response_data["id"],
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"max_tries": 0,
"note": "placeholder-note",
"operator": "PythonOperator",
@@ -588,6 +602,8 @@ class TestGetTaskInstance(TestTaskInstanceEndpoint):
"hostname": "",
"id": response_data["id"],
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"max_tries": 0,
"note": "placeholder-note",
"operator": "PythonOperator",
@@ -711,6 +727,8 @@ class TestGetMappedTaskInstance(TestTaskInstanceEndpoint):
"hostname": "",
"id": response_data["id"],
"map_index": map_index,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": map_index,
"max_tries": 0,
"note": "placeholder-note",
"operator": "PythonOperator",
@@ -3265,6 +3283,452 @@ class TestGetTaskInstanceTry(TestTaskInstanceEndpoint):
class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
+ @pytest.mark.parametrize("one_run", [False, True])
+ def test_broad_clear_uses_pinned_loop_after_latest_definition_removes_it(
+ self, test_client, dag_maker, session, one_run
+ ):
+ @task_group
+ def body():
+ MockOperator(task_id="member")
+
+ with dag_maker(dag_id="changed_loop", serialized=True):
+ loop = create_loop(body, max_iterations=2)
+ dr = dag_maker.create_dagrun()
+ root = session.scalar(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id))
+ for ti in dr.task_instances:
+ ti.state = State.SUCCESS
+ for loop_task in loop.iter_tasks():
+ session.add(
+ TaskInstance(
+ task=loop_task,
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=root.id,
+ region_index=1,
+ state=State.SUCCESS,
+ )
+ )
+ session.commit()
+ suffix_ids = {ti.id for ti in dr.get_task_instances(session=session) if ti.region_index == 1}
+ dag_id, run_id, gate_id = dr.dag_id, dr.run_id, loop.gate_task_id
+ with dag_maker(dag_id=dag_id, serialized=True, session=session):
+ MockOperator(task_id="replacement")
+ session.commit()
+
+ response = test_client.post(
+ f"/dags/{dag_id}/clearTaskInstances",
+ json={
+ "task_ids": [gate_id],
+ "dry_run": False,
+ "only_failed": False,
+ **({"dag_run_id": run_id} if one_run else {}),
+ },
+ )
+
+ assert response.status_code == 200, response.text
+ session.expire_all()
+ assert len(suffix_ids) == 2
+ for identity in suffix_ids:
+ archived = session.get(TaskInstance, identity)
+ assert (archived.working_set, archived.archived_reason) == (None, "superseded")
+ assert (
+ session.scalar(
+ select(func.count()).select_from(DynamicRegion).where(DynamicRegion.dag_id == dag_id)
+ )
+ == 2
+ )
+
+ @pytest.mark.parametrize("exact", [False, True])
+ def test_task_name_clear_refreshes_execution_after_waiting_for_dagrun_lock(
+ self, test_client, dag_maker, session, mocker, exact
+ ):
+ with dag_maker(serialized=True):
+ MockOperator(task_id="task")
+ dr = dag_maker.create_dagrun()
+ ti = dr.get_task_instance("task", session=session)
+ ti.state = State.SUCCESS
+ session.commit()
+ run_id, dag_id, original_id = dr.run_id, dr.dag_id, ti.id
+ execute = Session._execute_internal
+ bind = session.get_bind()
+ replacement_id = None
+
+ def overlap(request_session, statement, *args, **kwargs):
+ nonlocal replacement_id
+ run_lock = (
+ isinstance(statement, Select)
+ and statement._for_update_arg is not None
+ and any(getattr(table, "name", None) == "dag_run" for table in statement.get_final_froms())
+ )
+ if run_lock and replacement_id is None:
+ replacement_id = original_id
+ with Session(bind=bind) as other:
+ successor = clear_task_instances([other.get(TaskInstance, original_id)], session=other)[0]
+ successor.state = State.SUCCESS
+ other.commit()
+ replacement_id = successor.id
+ return execute(request_session, statement, *args, **kwargs)
+
+ mocker.patch.object(Session, "_execute_internal", autospec=True, side_effect=overlap)
+ selection = {"task_instance_ids": [str(original_id)]} if exact else {"task_ids": ["task"]}
+ response = test_client.post(
+ f"/dags/{dag_id}/clearTaskInstances",
+ json={"dag_run_id": run_id, **selection, "dry_run": False, "only_failed": False},
+ )
+
+ assert response.status_code == (409 if exact else 200), response.text
+ if not exact:
+ assert response.json()["total_entries"] == 1
+ session.expire_all()
+ assert session.get(TaskInstance, original_id).working_set is None
+ assert (session.get(TaskInstance, replacement_id).working_set is True) is exact
+
+ def test_exact_clear_rejects_missing_execution_without_changing_run(
+ self, test_client, dag_maker, session
+ ):
+ with dag_maker(serialized=True):
+ MockOperator(task_id="task")
+ dr = dag_maker.create_dagrun()
+ session.commit()
+ before = {ti.id for ti in dr.get_task_instances(session=session)}
+
+ response = test_client.post(
+ f"/dags/{dr.dag_id}/clearTaskInstances",
+ json={
+ "dag_run_id": dr.run_id,
+ "task_instance_ids": [str(uuid7())],
+ "dry_run": False,
+ "only_failed": False,
+ },
+ )
+
+ assert response.status_code == 404
+ session.expire_all()
+ assert dr.clear_number == 0
+ assert {ti.id for ti in dr.get_task_instances(session=session)} == before
+
+ @pytest.mark.parametrize(
+ "selection",
+ [
+ {"task_instance_ids": []},
+ {"task_instance_ids": ["00000000-0000-0000-0000-000000000001"]},
+ {
+ "task_instance_ids": ["00000000-0000-0000-0000-000000000001"],
+ "dag_run_id": "run",
+ "task_ids": ["task"],
+ },
+ {
+ "task_instance_ids": ["00000000-0000-0000-0000-000000000001"],
+ "dag_run_id": "run",
+ "task_group_id": "group",
+ },
+ {
+ "task_instance_ids": ["00000000-0000-0000-0000-000000000001"],
+ "dag_run_id": "run",
+ "include_future": True,
+ },
+ {
+ "task_instance_ids": ["00000000-0000-0000-0000-000000000001"],
+ "dag_run_id": "run",
+ "include_past": True,
+ },
+ {"dag_run_id": "run", "whole_expansion_ids": ["00000000-0000-0000-0000-000000000001"]},
+ ],
+ )
+ def test_exact_clear_rejects_conflicting_scope(self, test_client, selection):
+ response = test_client.post("/dags/example_python_operator/clearTaskInstances", json=selection)
+ assert response.status_code == 422
+
+ @pytest.mark.parametrize("downstream", [False, True])
+ @pytest.mark.parametrize("later", [False, True])
+ @pytest.mark.parametrize("dry_run", [False, True])
+ @pytest.mark.parametrize("only_failed", [False, True])
+ def test_exact_loop_clear_keeps_iteration_scope(
+ self, test_client, dag_maker, session, downstream, later, dry_run, only_failed
+ ):
+ @task_group
+ def body():
+ MockOperator(task_id="member")
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ loop >> MockOperator(task_id="outside")
+ dr = dag_maker.create_dagrun()
+ root = session.scalar(
+ select(DynamicRegion).where(
+ DynamicRegion.dag_id == dr.dag_id, DynamicRegion.node_id == loop.group_id
+ )
+ )
+ for index in (1, 2):
+ for loop_task in loop.iter_tasks():
+ session.add(
+ TaskInstance(
+ task=loop_task,
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=root.id,
+ region_index=index,
+ state=State.SUCCESS,
+ )
+ )
+ session.commit()
+ tis = list(dr.get_task_instances(session=session))
+ seed = next(ti for ti in tis if ti.task_id == "body.member" and ti.region_index == 1)
+ if only_failed:
+ seed.state = State.FAILED
+ session.commit()
+ before = {(ti.id, ti.state, ti.try_number) for ti in tis}
+ coordinates = {ti.id: (ti.task_id, str(ti.region_id), ti.region_index) for ti in tis}
+
+ response = test_client.post(
+ f"/dags/{dr.dag_id}/clearTaskInstances",
+ json={
+ "dag_run_id": dr.run_id,
+ "dry_run": dry_run,
+ "task_instance_ids": [str(seed.id)],
+ "only_failed": only_failed,
+ "include_downstream": downstream,
+ "include_later_loop_iterations": later,
+ },
+ )
+
+ assert response.status_code == 200, response.text
+ expected = {seed.id}
+ if downstream and not only_failed:
+ expected.update(
+ ti.id
+ for ti in tis
+ if ti.task_id == "outside"
+ or (ti.region_index == 1 and ti.task_id == loop.gate_task_id)
+ or (later and ti.region_id == root.id and ti.region_index > 1)
+ )
+ assert {
+ (row["task_id"], row["region_id"], row["region_index"])
+ for row in response.json()["task_instances"]
+ } == {coordinates[value] for value in expected}
+ session.expire_all()
+ if dry_run:
+ assert {
+ (ti.id, ti.state, ti.try_number) for ti in dr.get_task_instances(session=session)
+ } == before
+ else:
+ for identity, state, try_number in before:
+ if identity not in expected:
+ retained = session.get(TaskInstance, identity)
+ assert (retained.state, retained.try_number) == (state, try_number)
+ elif state in (State.SUCCESS, State.FAILED):
+ historical = session.get(TaskInstance, identity)
+ assert historical.working_set is None
+ archived = coordinates[identity][2] > 1 and downstream and later
+ assert historical.archived_reason == ("superseded" if archived else "retry")
+
+ def test_loop_clear_reports_conflict_when_a_worker_archived_the_execution_first(
+ self, test_client, dag_maker, session, mocker
+ ):
+ @task_group
+ def body():
+ MockOperator(task_id="member")
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ root = session.scalar(
+ select(DynamicRegion).where(
+ DynamicRegion.dag_id == dr.dag_id, DynamicRegion.node_id == loop.group_id
+ )
+ )
+ for loop_task in loop.iter_tasks():
+ session.add(
+ TaskInstance(
+ task=loop_task,
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=root.id,
+ region_index=1,
+ state=State.SUCCESS,
+ )
+ )
+ session.commit()
+ gate = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == loop.gate_task_id and ti.region_index == 0
+ )
+ mocker.patch.object(
+ TaskInstance,
+ "archive",
+ autospec=True,
+ side_effect=ValueError("An archived task instance cannot be archived again"),
+ )
+
+ response = test_client.post(
+ f"/dags/{dr.dag_id}/clearTaskInstances",
+ json={
+ "dag_run_id": dr.run_id,
+ "dry_run": False,
+ "only_failed": False,
+ "task_instance_ids": [str(gate.id)],
+ "include_later_loop_iterations": True,
+ },
+ )
+
+ assert response.status_code == 409, response.text
+
+ @pytest.mark.parametrize("downstream", [False, True])
+ @pytest.mark.parametrize("later", [False, True])
+ @pytest.mark.parametrize("seed_task", ["body.improve", "body.evaluate", "gate"])
+ @pytest.mark.parametrize("exact", [False, True])
+ @pytest.mark.parametrize("only_failed", [False, True])
+ def test_loop_clear_dry_run_reports_what_the_real_clear_replaces(
+ self, test_client, dag_maker, session, downstream, later, seed_task, exact, only_failed
+ ):
+ @task_group
+ def body():
+ improve = MockOperator(task_id="improve")
+ improve >> MockOperator(task_id="evaluate")
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=4)
+ loop >> MockOperator(task_id="finished")
+ dr = dag_maker.create_dagrun()
+ root = session.scalar(
+ select(DynamicRegion).where(
+ DynamicRegion.dag_id == dr.dag_id, DynamicRegion.node_id == loop.group_id
+ )
+ )
+ for index in (1, 2, 3):
+ for loop_task in loop.iter_tasks():
+ session.add(
+ TaskInstance(
+ task=loop_task,
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=root.id,
+ region_index=index,
+ state=State.SUCCESS,
+ )
+ )
+ for ti in dr.get_task_instances(session=session):
+ ti.state = State.FAILED if only_failed and ti.region_index in (1, 2) else State.SUCCESS
+ session.commit()
+ task_id = loop.gate_task_id if seed_task == "gate" else seed_task
+
+ def build_payload():
+ seed = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.run_id == dr.run_id,
+ TaskInstance.task_id == task_id,
+ TaskInstance.region_index == 1,
+ TaskInstance.working_set.is_(True),
+ )
+ )
+ return {
+ "dag_run_id": dr.run_id,
+ "only_failed": only_failed,
+ "include_downstream": downstream,
+ "include_later_loop_iterations": later,
+ **({"task_instance_ids": [str(seed.id)]} if exact else {"task_ids": [task_id]}),
+ }
+
+ test_client.post(f"/dags/{dr.dag_id}/clearTaskInstances", json={**build_payload(), "dry_run": False})
+ for ti in session.scalars(select(TaskInstance).where(TaskInstance.working_set.is_(True))):
+ ti.state = State.FAILED if only_failed and ti.region_index in (1, 2) else State.SUCCESS
+ session.commit()
+ payload = build_payload()
+ archived_before = set(
+ session.scalars(
+ select(TaskInstance.id)
+ .where(TaskInstance.working_set.is_(None))
+ .execution_options(include_all_attempts=True)
+ )
+ )
+ dry = test_client.post(f"/dags/{dr.dag_id}/clearTaskInstances", json={**payload, "dry_run": True})
+ real = test_client.post(f"/dags/{dr.dag_id}/clearTaskInstances", json={**payload, "dry_run": False})
+
+ assert dry.status_code == 200, dry.text
+ assert real.status_code == 200, real.text
+ session.expire_all()
+ archived_ids = {
+ ti.id
+ for ti in session.scalars(
+ select(TaskInstance)
+ .where(
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.working_set.is_(None),
+ )
+ .execution_options(include_all_attempts=True)
+ )
+ } - archived_before
+ reported = {
+ (row["task_id"], row["region_id"], row["region_index"]) for row in dry.json()["task_instances"]
+ }
+ replaced = {
+ (ti.task_id, str(ti.region_id), ti.region_index)
+ for ti in session.scalars(select(TaskInstance).where(TaskInstance.id.in_(archived_ids)))
+ }
+ assert reported == replaced
+ assert {
+ (row["task_id"], row["region_id"], row["region_index"]) for row in real.json()["task_instances"]
+ } >= reported
+
+ @pytest.mark.parametrize("new_note", [None, "Reason for clearing", ""])
+ @pytest.mark.parametrize("whole", [False, True])
+ @pytest.mark.parametrize("exact", [False, True])
+ def test_legacy_whole_clear_retains_affected_execution_and_note(
+ self, test_client, dag_maker, session, new_note, whole, exact
+ ):
+ with dag_maker(dag_id="legacy_clear_response", serialized=True):
+ MockOperator.partial(task_id="mapped").expand(arg2=[1])
+ dr = dag_maker.create_dagrun()
+ session.execute(delete(TaskInstance).where(TaskInstance.dag_id == dr.dag_id))
+ session.execute(delete(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id))
+ ti = TaskInstance(
+ dag_maker.serialized_dag.get_task("mapped"),
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_index=0,
+ state=State.SUCCESS,
+ )
+ ti.note = ("Original execution note", "test")
+ session.add(ti)
+ session.commit()
+ old_id = ti.id
+ selection = (
+ {"task_instance_ids": [str(old_id)], "whole_expansion_ids": [str(old_id)] if whole else []}
+ if exact
+ else {"task_ids": ["mapped"] if whole else [["mapped", 0]]}
+ )
+
+ response = test_client.post(
+ f"/dags/{dr.dag_id}/clearTaskInstances",
+ json={
+ "dry_run": False,
+ "reset_dag_runs": False,
+ "only_failed": False,
+ "dag_run_id": dr.run_id,
+ **selection,
+ "note": new_note,
+ },
+ )
+
+ assert response.status_code == 200, response.text
+ expected_note = "Original execution note" if new_note is None else new_note or None
+ body = response.json()
+ assert body["total_entries"] == 1
+ assert (body["task_instances"][0]["id"] == str(old_id)) is whole
+ assert body["task_instances"][0]["note"] == (expected_note if whole else new_note or None)
+ session.expire_all()
+ history = session.get(TaskInstance, old_id)
+ assert history.working_set is None
+ assert history.archived_reason == ("superseded" if whole else "retry")
+ assert history.note == (expected_note if whole else "Original execution note")
+ assert (
+ session.scalar(
+ select(func.count()).select_from(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)
+ )
+ == whole
+ )
+
@pytest.mark.parametrize(
("main_dag", "task_instances", "request_dag", "payload", "expected_ti"),
[
@@ -3781,7 +4245,12 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
# dag (3rd argument) is a different session object. Manually asserting that the dag_id
# is the same.
mock_clearti.assert_called_once_with(
- [], mock.ANY, DagRunState.QUEUED, prevent_running_task=False, run_on_latest_version=False
+ [],
+ mock.ANY,
+ DagRunState.QUEUED,
+ prevent_running_task=False,
+ run_on_latest_version=False,
+ whole_task_keys=set(),
)
def test_clear_taskinstance_is_called_with_invalid_task_ids(self, test_client, session):
@@ -3956,6 +4425,8 @@ class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
"id": response_data["task_instances"][0]["id"],
"logical_date": response_logical_date,
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"max_tries": 0,
"note": "placeholder-note",
"operator": "PythonOperator",
@@ -4940,6 +5411,474 @@ class TestGetTaskInstanceTries(TestTaskInstanceEndpoint):
}
+class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
+ @pytest.fixture
+ def loop_instances(self, dag_maker, session):
+ @task_group
+ def body():
+ MockOperator(task_id="first") >> MockOperator(task_id="last")
+
+ with dag_maker("manual-loop", serialized=True):
+ loop = create_loop(body, max_iterations=3, until=lambda loop: True)
+ dr = dag_maker.create_dagrun(run_id="manual-run")
+ root = session.scalar(select(DynamicRegion).where(DynamicRegion.node_id == loop.group_id))
+ for index in (1, 2):
+ for operator in loop.iter_tasks():
+ session.add(
+ TaskInstance(
+ task=operator,
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=root.id,
+ region_index=index,
+ )
+ )
+ session.flush()
+ tis = {(ti.task_id, ti.region_index): ti for ti in dr.get_task_instances(session=session)}
+ for ti in tis.values():
+ ti.set_state(State.FAILED, session=session)
+ session.commit()
+ return dr, loop, root, tis
+
+ @pytest.mark.parametrize("keep_task_state", [False, True])
+ def test_clear_discards_task_state_of_the_selected_loop_iteration_only(
+ self, test_client, session, loop_instances, keep_task_state
+ ):
+ dr, loop, root, tis = loop_instances
+ backend = MetastoreBackend()
+ for index in (1, 2):
+ backend.set(
+ TaskScope(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ task_id="body.first",
+ region_id=root.id,
+ region_index=index,
+ ),
+ "job_id",
+ "external-job",
+ session=session,
+ )
+ session.commit()
+
+ response = test_client.post(
+ f"/dags/{dr.dag_id}/clearTaskInstances",
+ json={
+ "dry_run": False,
+ "dag_run_id": dr.run_id,
+ "task_instance_ids": [str(tis["body.first", 1].id)],
+ "keep_task_state": keep_task_state,
+ },
+ )
+
+ assert response.status_code == 200, response.text
+ session.expire_all()
+ assert set(
+ session.scalars(
+ select(TaskStateStoreModel.region_index).where(
+ TaskStateStoreModel.dag_id == dr.dag_id,
+ TaskStateStoreModel.task_id == "body.first",
+ )
+ )
+ ) == ({1, 2} if keep_task_state else {2})
+
+ @pytest.mark.parametrize("dry_run", [False, True])
+ def test_exact_loop_mark_does_not_change_other_iterations(
+ self, test_client, session, loop_instances, dry_run
+ ):
+ dr, loop, root, tis = loop_instances
+ url = f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/body.first"
+
+ response = test_client.patch(
+ url + ("/dry_run" if dry_run else ""),
+ json={"new_state": "success", "region_id": str(root.id), "region_index": 1},
+ )
+
+ assert response.status_code == 200, response.text
+ assert [
+ (item["task_id"], item["map_index"], item["region_index"])
+ for item in response.json()["task_instances"]
+ ] == [("body.first", -1, 1)]
+ session.expire_all()
+ for index in (0, 2):
+ assert tis["body.first", index].state == State.FAILED
+ assert tis["body.last", index].state == State.FAILED
+ assert tis["body.first", 1].state == (State.FAILED if dry_run else State.SUCCESS)
+ current_last = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.run_id == dr.run_id,
+ TaskInstance.task_id == "body.last",
+ TaskInstance.region_id == root.id,
+ TaskInstance.region_index == 1,
+ TaskInstance.working_set.is_(True),
+ )
+ )
+ assert current_last.state == (State.FAILED if dry_run else None)
+
+ def test_loop_mark_without_coordinates_is_ambiguous(self, test_client, loop_instances):
+ dr, loop, root, tis = loop_instances
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/body.first",
+ json={"new_state": "success"},
+ )
+
+ assert response.status_code == 409, response.text
+
+ @pytest.mark.parametrize("coordinates_in_query", [False, True])
+ def test_exact_loop_note_changes_only_selected_execution(
+ self, test_client, session, loop_instances, coordinates_in_query
+ ):
+ dr, loop, root, tis = loop_instances
+ coordinates = {"region_id": str(root.id), "region_index": 1}
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/body.first",
+ json={"note": "selected iteration", **({} if coordinates_in_query else coordinates)},
+ params=coordinates if coordinates_in_query else {},
+ )
+
+ assert response.status_code == 200, response.text
+ session.expire_all()
+ assert tis["body.first", 1].note == "selected iteration"
+ assert tis["body.first", 0].note is None
+ assert tis["body.first", 2].note is None
+
+ def test_manual_gate_success_after_rewind_stops_without_consuming_continue(
+ self, test_client, session, loop_instances
+ ):
+ dr, loop, root, tis = loop_instances
+ clear_loop_task_instances([tis[loop.gate_task_id, 0]], downstream=False, session=session)
+ gate = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.task_id == loop.gate_task_id,
+ TaskInstance.region_index == 0,
+ TaskInstance.working_set.is_(True),
+ )
+ ).one()
+ gate_id = gate.id
+ session.add(XComModelV2(task_instance_id=gate_id, key=LOOP_DECISION_KEY, value="continue"))
+ session.commit()
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/{gate.task_id}",
+ json={"new_state": "success", "region_id": str(root.id), "region_index": 0},
+ )
+
+ assert response.status_code == 200, response.text
+ session.expire_all()
+ assert session.get(TaskInstance, gate_id).state == State.SUCCESS
+ assert all(ti.region_index == 0 for ti in dr.get_task_instances(session=session))
+ assert not session.scalar(select(XComModelV2).where(XComModelV2.task_instance_id == gate_id))
+
+ @pytest.mark.parametrize("outcome", ["success", "failed", "skipped"])
+ def test_manual_terminal_settles_superseded_running_execution(
+ self, test_client, session, loop_instances, outcome
+ ):
+ dr, loop, root, tis = loop_instances
+ archiving = tis["body.first", 2]
+ archiving.state = State.RUNNING
+ session.flush()
+ archiving_id = archiving.id
+ clear_loop_task_instances([tis[loop.gate_task_id, 0]], downstream=False, session=session)
+ assert archiving.state == State.RESTARTING
+ session.commit()
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/{archiving.task_id}",
+ json={"new_state": outcome, "region_id": str(root.id), "region_index": 2},
+ )
+
+ assert response.status_code == 200, response.text
+ assert response.json()["task_instances"][0]["state"] == outcome
+ session.expire_all()
+ history = session.get(TaskInstance, archiving_id)
+ assert history.working_set is None
+ assert history.state == outcome
+ assert history.archived_reason == "superseded"
+ assert history.end_date is not None
+
+ @pytest.mark.parametrize(
+ ("body", "params"),
+ [
+ ({"region_index": 1}, {}),
+ ({}, {"region_index": 1}),
+ (
+ {"region_index": 1, "region_id": "00000000-0000-0000-0000-000000000000"},
+ {"region_index": 2, "region_id": "00000000-0000-0000-0000-000000000000"},
+ ),
+ (
+ {
+ "region_index": 1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "include_past": True,
+ },
+ {},
+ ),
+ ],
+ )
+ def test_invalid_coordinate_control_is_rejected(self, test_client, loop_instances, body, params):
+ dr, loop, root, tis = loop_instances
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/body.first",
+ json={"new_state": "success", **body},
+ params=params,
+ )
+
+ assert response.status_code == 400, response.text
+
+ @pytest.mark.parametrize("action", ["update", "delete"])
+ def test_bulk_exact_passes_have_distinct_execution_results(
+ self, test_client, session, loop_instances, action
+ ):
+ dr, loop, root, tis = loop_instances
+ ids = {str(tis["body.first", index].id) for index in (0, 2)}
+ entities = [
+ {
+ "task_id": "body.first",
+ "map_index": -1,
+ "region_id": str(root.id),
+ "region_index": index,
+ **({"new_state": "success"} if action == "update" else {}),
+ }
+ for index in (0, 2)
+ ]
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances",
+ json={"actions": [{"action": action, "entities": entities}]},
+ )
+
+ assert response.status_code == 200, response.text
+ assert not response.json()[action]["errors"]
+ assert set(response.json()[action]["success"]) == ids
+ session.expire_all()
+ assert tis["body.first", 1].state == State.FAILED
+ live = {str(ti.id): ti.state for ti in dr.get_task_instances(session=session)}
+ if action == "delete":
+ assert not ids & live.keys()
+ else:
+ assert all(live[ti_id] == State.SUCCESS for ti_id in ids)
+
+ @pytest.mark.parametrize("dry_run", [False, True])
+ def test_group_coordinates_select_one_loop_pass(self, test_client, session, loop_instances, dry_run):
+ dr, loop, root, tis = loop_instances
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskGroupInstances/body"
+ + ("/dry_run" if dry_run else ""),
+ json={"new_state": "success", "region_id": str(root.id), "region_index": 1},
+ )
+
+ assert response.status_code == 200, response.text
+ assert {item["id"] for item in response.json()["task_instances"]} == {
+ str(ti.id) for (task_id, index), ti in tis.items() if index == 1
+ }
+ session.expire_all()
+ for (_task_id, index), ti in tis.items():
+ assert ti.state == (State.SUCCESS if index == 1 and not dry_run else State.FAILED)
+
+ @pytest.mark.parametrize("action", ["update", "delete"])
+ @pytest.mark.parametrize("map_index", [None, -1])
+ def test_bulk_loop_request_requires_coordinates(
+ self, test_client, session, loop_instances, action, map_index
+ ):
+ dr, loop, root, tis = loop_instances
+ before = {ti.id: ti.state for ti in tis.values()}
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances",
+ json={
+ "actions": [
+ {
+ "action": action,
+ "entities": [
+ {
+ "task_id": "body.first",
+ "map_index": map_index,
+ **({"new_state": "success"} if action == "update" else {}),
+ }
+ ],
+ }
+ ]
+ },
+ )
+
+ assert response.status_code == 200, response.text
+ assert response.json()[action]["errors"][0]["status_code"] == 409
+ assert not response.json()[action]["success"]
+ session.expire_all()
+ assert {ti.id: ti.state for ti in dr.get_task_instances(session=session)} == before
+
+ def test_group_pass_includes_mapped_descendants(self, test_client, dag_maker, session):
+ @task_group
+ def body():
+ MockOperator.partial(task_id="mapped").expand(arg2=[1, 2]) >> MockOperator(task_id="last")
+
+ with dag_maker("mapped-manual-loop", serialized=True):
+ loop = create_loop(body, max_iterations=2)
+ loop >> MockOperator(task_id="outside")
+ dr = dag_maker.create_dagrun(run_id="manual-run")
+ root = session.scalar(select(DynamicRegion).where(DynamicRegion.node_id == loop.group_id))
+ session.commit()
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskGroupInstances/body",
+ json={"new_state": "success", "region_id": str(root.id), "region_index": 0},
+ )
+
+ assert response.status_code == 200, response.text
+ values = response.json()["task_instances"]
+ assert {(ti["task_id"], ti["map_index"]) for ti in values} == {
+ ("body.mapped", 0),
+ ("body.mapped", 1),
+ ("body.last", -1),
+ (loop.gate_task_id, -1),
+ }
+ session.expire_all()
+ assert all(
+ ti.state == (None if ti.task_id == "outside" else State.SUCCESS)
+ for ti in dr.get_task_instances(session=session)
+ )
+
+ @conf_vars({("state_store", "clear_on_success"): "True"})
+ def test_regional_success_preserves_other_state_scope_and_notifies_listener(
+ self, test_client, session, loop_instances, listener_manager
+ ):
+ dr, loop, root, tis = loop_instances
+ backend = MetastoreBackend()
+ scopes = [
+ TaskScope(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ task_id="body.first",
+ region_id=root.id,
+ region_index=index,
+ )
+ for index in (0, 1)
+ ]
+ for scope in scopes:
+ backend.set(scope, "job_id", "external-job", session=session)
+ session.commit()
+ listener = ClassBasedListener()
+ listener_manager(listener)
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/body.first",
+ json={
+ "new_state": "success",
+ "note": "completed manually",
+ "region_id": str(root.id),
+ "region_index": 1,
+ },
+ )
+
+ assert response.status_code == 200, response.text
+ assert listener.state == [TaskInstanceState.SUCCESS]
+ assert listener.ti_note_at_listener == "completed manually"
+ assert set(
+ session.scalars(
+ select(TaskStateStoreModel.region_index).where(
+ TaskStateStoreModel.dag_id == dr.dag_id,
+ TaskStateStoreModel.task_id == "body.first",
+ )
+ )
+ ) == {0}
+
+ @pytest.mark.parametrize("dry_run", [False, True])
+ def test_regional_downstream_mark_uses_same_scope_as_preview(
+ self, test_client, session, loop_instances, dry_run
+ ):
+ dr, loop, root, tis = loop_instances
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskInstances/body.first"
+ + ("/dry_run" if dry_run else ""),
+ json={
+ "new_state": "success",
+ "include_downstream": True,
+ "region_id": str(root.id),
+ "region_index": 1,
+ },
+ )
+
+ assert response.status_code == 200, response.text
+ assert {item["id"] for item in response.json()["task_instances"]} == {
+ str(ti.id) for (_task_id, index), ti in tis.items() if index == 1
+ }
+ session.expire_all()
+ assert all(
+ ti.state == (State.SUCCESS if index == 1 and not dry_run else State.FAILED)
+ for (_task_id, index), ti in tis.items()
+ )
+
+ @pytest.mark.parametrize("remove_group", [False, True])
+ def test_regional_group_uses_execution_version_after_definition_changes(
+ self, test_client, session, dag_maker, loop_instances, remove_group
+ ):
+ dr, loop, root, tis = loop_instances
+
+ @task_group(group_id="replacement" if remove_group else "body")
+ def changed():
+ MockOperator(task_id="first")
+
+ with dag_maker(dr.dag_id, serialized=True, session=session):
+ create_loop(changed, max_iterations=3, until=lambda loop: True)
+ session.commit()
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskGroupInstances/body",
+ json={"new_state": "success", "region_id": str(root.id), "region_index": 1},
+ )
+
+ assert response.status_code == 200, response.text
+ assert {item["id"] for item in response.json()["task_instances"]} == {
+ str(ti.id) for (_task_id, index), ti in tis.items() if index == 1
+ }
+
+ def test_group_terminal_response_retains_executions_archived_by_manual_settlement(
+ self, test_client, session, loop_instances
+ ):
+ dr, loop, root, tis = loop_instances
+ draining = [tis["body.first", 2], tis[loop.gate_task_id, 2]]
+ for ti in draining:
+ ti.state = State.RUNNING
+ session.flush()
+ ids = {str(ti.id) for ti in draining}
+ clear_loop_task_instances([tis[loop.gate_task_id, 0]], downstream=False, session=session)
+ session.commit()
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskGroupInstances/body",
+ json={"new_state": "skipped", "region_id": str(root.id), "region_index": 2},
+ )
+
+ assert response.status_code == 200, response.text
+ assert {item["id"] for item in response.json()["task_instances"]} == ids
+ assert all(item["state"] == "skipped" for item in response.json()["task_instances"])
+
+ def test_unscoped_group_rejects_pinned_loop_when_latest_group_is_plain(
+ self, test_client, session, dag_maker, loop_instances
+ ):
+ dr, loop, root, tis = loop_instances
+ before = {ti.id: ti.state for ti in tis.values()}
+ with dag_maker(dr.dag_id, serialized=True, session=session):
+ with TaskGroup(group_id="body"):
+ MockOperator(task_id="first") >> MockOperator(task_id="last")
+ session.commit()
+
+ response = test_client.patch(
+ f"/dags/{dr.dag_id}/dagRuns/{dr.run_id}/taskGroupInstances/body",
+ json={"new_state": "success"},
+ )
+
+ assert response.status_code == 409, response.text
+ session.expire_all()
+ assert {ti.id: ti.state for ti in dr.get_task_instances(session=session)} == before
+
+
class TestPatchTaskInstance(TestTaskInstanceEndpoint):
ENDPOINT_URL = "/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context"
NEW_STATE = "failed"
@@ -5040,6 +5979,8 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
"hostname": "",
"id": response_data["task_instances"][0]["id"],
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"max_tries": 0,
"note": "placeholder-note",
"operator": "PythonOperator",
@@ -5319,6 +6260,8 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
"hostname": "",
"id": mock.ANY,
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"max_tries": 0,
"note": "placeholder-note",
"operator": "PythonOperator",
@@ -5458,6 +6401,8 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
"executor_config": "{}",
"hostname": "",
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"max_tries": 0,
"note": new_note_value,
"operator": "PythonOperator",
@@ -5524,6 +6469,8 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
"executor_config": "{}",
"hostname": "",
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"max_tries": 0,
"note": new_note_value,
"operator": "PythonOperator",
@@ -5623,6 +6570,8 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
"executor_config": "{}",
"hostname": "",
"map_index": map_index,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": map_index,
"max_tries": 0,
"note": new_note_value,
"operator": "PythonOperator",
@@ -5710,6 +6659,8 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
"executor_config": "{}",
"hostname": "",
"map_index": map_index,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": map_index,
"max_tries": 0,
"note": new_note_value,
"operator": "PythonOperator",
@@ -5909,6 +6860,8 @@ class TestPatchTaskInstanceDryRun(TestTaskInstanceEndpoint):
"hostname": "",
"id": response_data["task_instances"][0]["id"],
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"max_tries": 0,
"note": "placeholder-note",
"operator": "PythonOperator",
@@ -6200,6 +7153,8 @@ class TestPatchTaskInstanceDryRun(TestTaskInstanceEndpoint):
"hostname": "",
"id": mock.ANY,
"map_index": -1,
+ "region_id": "00000000-0000-0000-0000-000000000000",
+ "region_index": -1,
"max_tries": 0,
"note": "placeholder-note",
"operator": "PythonOperator",
@@ -7330,7 +8285,7 @@ class TestBulkTaskInstances(TestTaskInstanceEndpoint):
def test_bulk_delete_query_count_scales_linearly_with_task_count(self, test_client, session, task_count):
# Each extra task instance adds one coordinate DELETE, with no per-instance re-SELECT.
QUERIES_PER_TASK_INSTANCE = 1
- BASE_QUERY_COUNT = 2
+ BASE_QUERY_COUNT = 3
self.create_task_instances(
session,
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 aa7ec5f706f..794f423402d 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
@@ -49,7 +49,10 @@ 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.taskinstance import (
+ TISuccessStatePayload,
+ TITerminalStatePayload,
+)
from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
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
@@ -62,6 +65,7 @@ from airflow.models.dag import DagModel
from airflow.models.dagbag import DBDagBag
from airflow.models.dynamic_region import DynamicRegion
from airflow.models.log import Log
+from airflow.models.loop_clear import clear_loop_task_instances
from airflow.models.task_coordinates import TaskCoordinateResolver
from airflow.models.task_state_store import TaskStateStoreModel
from airflow.models.taskinstance import TaskInstance, clear_task_instances
@@ -178,6 +182,98 @@ def test_id_matches_sub_claim(client, session, create_task_instance):
validator.avalidated_claims.assert_awaited()
+@pytest.mark.parametrize("ack_interleaving", ["before_reconciliation", "during_reconciliation"])
+def test_legacy_expansion_waits_for_all_cleared_workers_before_replacement(
+ client, dag_maker, session, ack_interleaving
+):
+ with dag_maker(serialized=True):
+ PythonOperator.partial(task_id="mapped", python_callable=str).expand(op_args=[[1], [2]])
+ dr = dag_maker.create_dagrun()
+ session.execute(delete(TaskInstance).where(TaskInstance.dag_id == dr.dag_id))
+ session.execute(delete(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id))
+ old = [
+ TaskInstance(
+ dag_maker.serialized_dag.get_task("mapped"),
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_index=index,
+ state=State.QUEUED,
+ )
+ for index in range(2)
+ ]
+ session.add_all(old)
+ session.commit()
+ old_ids = [ti.id for ti in old]
+ for ti_id in old_ids:
+ response = client.patch(
+ f"/execution/task-instances/{ti_id}/run",
+ json={
+ "state": "running",
+ "hostname": "original",
+ "unixname": "airflow",
+ "pid": 123,
+ "start_date": DEFAULT_START_DATE.isoformat(),
+ },
+ )
+ assert response.status_code == 200, response.text
+
+ session.expire_all()
+ for _ in range(2):
+ dag_maker.serialized_dag.clear(task_ids=["mapped"], run_id=dr.run_id, session=session)
+ dr.dag = dag_maker.serialized_dag
+ dr.verify_integrity(dag_version_id=dr.created_dag_version_id, session=session)
+ session.commit()
+ session.expire_all()
+ assert all(session.get(TaskInstance, ti_id).state == State.RESTARTING for ti_id in old_ids)
+ region = session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)).one()
+ assert not session.scalars(select(TaskInstance).where(TaskInstance.region_id == region.id)).all()
+
+ payload = {
+ "state": "server_terminated",
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "hostname": "original",
+ "pid": 123,
+ }
+ for index, ti_id in enumerate(old_ids):
+ if index == 1 and ack_interleaving == "during_reconciliation":
+ reconcile = dr._reconcile_legacy_expansions
+
+ def acknowledge_then_reconcile(*, session, ti_id=ti_id, reconcile=reconcile):
+ response = client.patch(f"/execution/task-instances/{ti_id}/state", json=payload)
+ assert response.status_code == 204, response.text
+ return reconcile(session=session)
+
+ with mock.patch.object(
+ dr, "_reconcile_legacy_expansions", autospec=True, side_effect=acknowledge_then_reconcile
+ ):
+ dr.update_state(session=session)
+ assert dr.state == DagRunState.RUNNING
+ else:
+ response = client.patch(f"/execution/task-instances/{ti_id}/state", json=payload)
+ assert response.status_code == 204, response.text
+ session.expire_all()
+ archived = session.get(TaskInstance, ti_id)
+ assert (archived.working_set, archived.archived_reason) == (None, "superseded")
+ assert archived.region_id.int == 0
+ if index == 0:
+ dr.dag = dag_maker.serialized_dag
+ dr.update_state(session=session)
+ assert dr.state == DagRunState.RUNNING
+ assert not session.scalars(select(TaskInstance).where(TaskInstance.region_id == region.id)).all()
+ session.commit()
+ assert client.patch(f"/execution/task-instances/{ti_id}/state", json=payload).status_code == 204
+
+ dr.dag = dag_maker.serialized_dag
+ dr.update_state(session=session)
+ assert dr.state == DagRunState.RUNNING
+ replacements = session.scalars(select(TaskInstance).where(TaskInstance.region_id == region.id)).all()
+ replacement_ids = {ti.id for ti in replacements}
+ assert len(replacement_ids) == 2
+ assert replacement_ids.isdisjoint(old_ids)
+ dr.update_state(session=session)
+ assert {ti.id for ti in dr.get_task_instances(session=session)} == replacement_ids
+
+
@pytest.fixture
def loop_reader_tis(dag_maker, session):
@task_group
@@ -1574,6 +1670,168 @@ class TestTIUpdateState:
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)}
+ @pytest.fixture
+ def started_loop_successor(self, client, session, running_loop_gate):
+ gate = running_loop_gate
+ dr = gate.dag_run
+ 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()
+ running = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id != gate.task_id and ti.region_index == 1
+ )
+ running.state = State.QUEUED
+ for ti in dr.get_task_instances(session=session):
+ if ti.task_id != gate.task_id and ti.region_index == 0:
+ ti.state = State.SUCCESS
+ session.commit()
+ running_id = running.id
+ response = client.patch(
+ f"/execution/task-instances/{running_id}/run",
+ json={
+ "state": "running",
+ "hostname": "archiving-worker",
+ "unixname": "airflow",
+ "pid": 123,
+ "start_date": DEFAULT_START_DATE.isoformat(),
+ },
+ )
+ assert response.status_code == 200, response.text
+ session.expire_all()
+ return gate, running_id
+
+ def test_loop_rewind_waits_for_worker_acknowledgement_before_gate_admission(
+ self, client, session, started_loop_successor, dag_maker
+ ):
+ gate, running_id = started_loop_successor
+ dr = gate.dag_run
+ running = session.get(TaskInstance, running_id)
+
+ for _ in range(2):
+ clear_loop_task_instances([session.get(TaskInstance, gate.id)], downstream=False, session=session)
+ session.commit()
+ session.expire_all()
+ assert running.id == running_id
+ assert running.state == State.RESTARTING
+ assert running.working_set is True
+ dr.dag = dag_maker.serialized_dag
+ decision = dr.task_instance_scheduling_decisions(session=session)
+ assert gate.id not in {ti.id for ti in decision.schedulable_tis}
+ dr.update_state(session=session)
+ assert dr.state == DagRunState.RUNNING
+ session.commit()
+ gate = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == gate.task_id and ti.region_index == gate.region_index
+ )
+
+ response = client.patch(
+ f"/execution/task-instances/{running_id}/state",
+ json={
+ "state": "server_terminated",
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "hostname": "archiving-worker",
+ "pid": 123,
+ },
+ )
+
+ assert response.status_code == 204, response.text
+ session.expire_all()
+ archived = session.get(TaskInstance, running_id)
+ assert (archived.working_set, archived.archived_reason) == (None, "superseded")
+ dr.dag = dag_maker.serialized_dag
+ decision = dr.task_instance_scheduling_decisions(session=session)
+ assert gate.id in {ti.id for ti in decision.schedulable_tis}
+ assert not any(ti.region_index == 1 for ti in dr.get_task_instances(session=session))
+
+ def test_restart_ack_locks_run_before_task(self, client, session, create_task_instance, mocker):
+ if session.bind.dialect.name == "sqlite":
+ pytest.skip("SQLite has no row locks")
+ ti = create_task_instance(state=State.QUEUED)
+ session.commit()
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/run",
+ json={
+ "state": "running",
+ "hostname": "worker",
+ "unixname": "airflow",
+ "pid": 123,
+ "start_date": DEFAULT_START_DATE.isoformat(),
+ },
+ )
+ assert response.status_code == 200, response.text
+ session.refresh(ti)
+ clear_task_instances([ti], session=session)
+ session.commit()
+ archiving_id = ti.id
+ dag_id, run_id = ti.dag_id, ti.run_id
+ ack_locked = Event()
+ clear_at_lock = Event()
+ execute = Session._execute_internal
+ bind = session.get_bind()
+
+ def coordinate(request_session, statement, *args, **kwargs):
+ locking = isinstance(statement, Select) and statement._for_update_arg is not None
+ run_lock = locking and any(
+ getattr(table, "name", None) == "dag_run" for table in statement.get_final_froms()
+ )
+ role = request_session.info.get("role")
+ if role == "clear" and run_lock:
+ clear_at_lock.set()
+ result = execute(request_session, statement, *args, **kwargs)
+ if role == "ack" and run_lock:
+ request_session.info["run_locked"] = True
+ elif role == "ack" and locking:
+ ack_locked.set()
+ assert request_session.info.get("run_locked")
+ assert clear_at_lock.wait(timeout=10)
+ return result
+
+ mocker.patch.object(Session, "_execute_internal", autospec=True, side_effect=coordinate)
+
+ def acknowledge():
+ with Session(bind=bind, info={"role": "ack"}) as other:
+ result = ti_update_state(
+ task_instance_id=archiving_id,
+ ti_patch_payload=TITerminalStatePayload(
+ state="server_terminated", end_date=DEFAULT_END_DATE, hostname="worker", pid=123
+ ),
+ session=other,
+ dag_bag=DBDagBag(),
+ )
+ other.commit()
+ return result.status_code
+
+ def clear_again():
+ assert ack_locked.wait(timeout=10)
+ with Session(bind=bind, info={"role": "clear"}) as other:
+ other.scalar(
+ select(DagRun).where(DagRun.dag_id == dag_id, DagRun.run_id == run_id).with_for_update()
+ )
+ selected = other.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.working_set.is_(True),
+ )
+ )
+ clear_task_instances([selected], session=other)
+ other.commit()
+
+ with ThreadPoolExecutor(max_workers=2) as pool:
+ ack = pool.submit(acknowledge)
+ repeated = pool.submit(clear_again)
+ assert ack.result(timeout=20) == 204
+ repeated.result(timeout=20)
+ session.expire_all()
+ assert session.get(TaskInstance, archiving_id).working_set is None
+
@conf_vars({("state_store", "clear_on_success"): "true"})
@pytest.mark.parametrize("mapped", [False, True])
@pytest.mark.parametrize(
@@ -2018,6 +2276,41 @@ class TestTIUpdateState:
session.refresh(ti)
assert ti.state == state
+ @pytest.mark.parametrize(
+ ("region_index", "locks_run"),
+ [(-1, False), (0, True)],
+ )
+ def test_stopped_report_locks_dag_run_only_for_regional_task(
+ self, client, session, create_task_instance, mocker, region_index, locks_run
+ ):
+ ti = create_task_instance(task_id="stopped_lock", state=State.SUCCESS)
+ ti.hostname = "worker"
+ ti.pid = 123
+ ti.region_index = region_index
+ session.commit()
+ execute = Session.execute
+ locked_tables = []
+
+ def record_locks(request_session, statement, *args, **kwargs):
+ if isinstance(statement, Select) and statement._for_update_arg is not None:
+ locked_tables.extend(getattr(table, "name", None) for table in statement.get_final_froms())
+ return execute(request_session, statement, *args, **kwargs)
+
+ mocker.patch.object(Session, "execute", autospec=True, side_effect=record_locks)
+
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/state",
+ json={
+ "state": "server_terminated",
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "hostname": "worker",
+ "pid": 123,
+ },
+ )
+
+ assert response.status_code == 204
+ assert ("dag_run" in locked_tables) == locks_run
+
@pytest.mark.parametrize("matching_worker", [True, False])
@pytest.mark.parametrize("missing_definition", [None, "dag", "task"])
@pytest.mark.parametrize("max_tries", [0, 7])
diff --git a/airflow-core/tests/unit/cli/commands/test_dag_command.py b/airflow-core/tests/unit/cli/commands/test_dag_command.py
index 8437e7410cb..36da45f1bfa 100644
--- a/airflow-core/tests/unit/cli/commands/test_dag_command.py
+++ b/airflow-core/tests/unit/cli/commands/test_dag_command.py
@@ -41,8 +41,9 @@ from airflow.dag_processing.processor import DagFileParsingResult, DagFileProces
from airflow.exceptions import AirflowException
from airflow.models import DagModel, DagRun
from airflow.models.dagbag import DBDagBag
+from airflow.models.loop_clear import clear_task_instances_for_runs
from airflow.models.serialized_dag import SerializedDagModel
-from airflow.models.taskinstance import TaskInstance, clear_task_instances
+from airflow.models.taskinstance import TaskInstance
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.triggers.temporal import DateTimeTrigger, TimeDeltaTrigger
from airflow.sdk import DAG, Asset, BaseOperator, CronPartitionTimetable, PartitionedAssetTimetable, task
@@ -2127,16 +2128,16 @@ class TestCliDagsClear:
def test_clears_each_matching_run_once_across_chunks(self, parser, chunk_size, expected_calls):
"""Every matching run is cleared exactly once, however run_ids split into chunks.
- clear_task_instances is called once per chunk (not once per run), every matching
+ clear_task_instances_for_runs is called once per chunk (not once per run), every matching
run is re-queued, and each run's clear_number advances by exactly 1 — proving a
run's TIs are never split across chunks.
"""
call_count = 0
- def counting_clear(tis, session, **kwargs):
+ def counting_clear(tis, **kwargs):
nonlocal call_count
call_count += 1
- return clear_task_instances(tis, session, **kwargs)
+ return clear_task_instances_for_runs(tis, **kwargs)
args = parser.parse_args(
[
@@ -2153,7 +2154,7 @@ class TestCliDagsClear:
with (
mock.patch.object(dag_command, "_RUN_CHUNK_SIZE", chunk_size),
mock.patch(
- "airflow.cli.commands.dag_command.clear_task_instances",
+ "airflow.cli.commands.dag_command.clear_task_instances_for_runs",
side_effect=counting_clear,
),
):
diff --git a/airflow-core/tests/unit/cli/commands/test_partition_command.py b/airflow-core/tests/unit/cli/commands/test_partition_command.py
index 7892c5bd74d..f707c3fae5b 100644
--- a/airflow-core/tests/unit/cli/commands/test_partition_command.py
+++ b/airflow-core/tests/unit/cli/commands/test_partition_command.py
@@ -435,7 +435,7 @@ class TestPartitionsClear:
with (
mock.patch("airflow.models.dagrun._TI_CHUNK_SIZE", ti_cap),
- mock.patch("airflow.models.dagrun.clear_task_instances", autospec=True) as mock_cti,
+ mock.patch("airflow.models.loop_clear.clear_task_instances_for_runs", autospec=True) as mock_cti,
):
partition_command.clear(
parser.parse_args(
@@ -486,7 +486,7 @@ class TestPartitionsClear:
with (
mock.patch("airflow.models.dagrun._TI_CHUNK_SIZE", ti_cap),
- mock.patch("airflow.models.dagrun.clear_task_instances", autospec=True) as mock_cti,
+ mock.patch("airflow.models.loop_clear.clear_task_instances_for_runs", autospec=True) as mock_cti,
):
partition_command.clear(
parser.parse_args(
@@ -538,7 +538,7 @@ class TestPartitionsClear:
with (
mock.patch("airflow.models.dagrun._TI_CHUNK_SIZE", ti_cap),
- mock.patch("airflow.models.dagrun.clear_task_instances", autospec=True) as mock_cti,
+ mock.patch("airflow.models.loop_clear.clear_task_instances_for_runs", autospec=True) as mock_cti,
):
partition_command.clear(
parser.parse_args(
@@ -611,7 +611,7 @@ class TestPartitionsClear:
with (
mock.patch("airflow.models.dagrun._TI_CHUNK_SIZE", ti_cap),
- mock.patch("airflow.models.dagrun.clear_task_instances", autospec=True) as mock_cti,
+ mock.patch("airflow.models.loop_clear.clear_task_instances_for_runs", autospec=True) as mock_cti,
):
partition_command.clear(
parser.parse_args(
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index dfadf308a39..75481fff5de 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -796,6 +796,78 @@ class TestSchedulerJob:
== 2
)
+ @pytest.mark.parametrize("locked_table", [DagRun, TaskInstance])
+ def test_executor_terminal_event_survives_locked_row(self, dag_maker, session, locked_table):
+ if session.bind.dialect.name == "sqlite":
+ pytest.skip("SQLite has no row locks")
+ with dag_maker(serialized=True):
+ EmptyOperator(task_id="task")
+ dr = dag_maker.create_dagrun()
+ ti = dr.get_task_instance("task", session=session)
+ ti.state = State.RUNNING
+ session.flush()
+ clear_task_instances([ti], session=session)
+ session.commit()
+ archiving_id = ti.id
+ executor = MockExecutor(do_update=False)
+ executor.event_buffer[TaskInstanceUuid(archiving_id)] = State.SUCCESS, "terminated"
+ runner = SchedulerJobRunner(Job(), executors=[executor])
+
+ def process():
+ with create_session(scoped=False) as other:
+ runner._process_executor_events(executor=executor, session=other)
+
+ with ThreadPoolExecutor(max_workers=1) as pool:
+ with create_session(scoped=False) as locker:
+ locked_id = dr.id if locked_table is DagRun else archiving_id
+ locker.scalar(select(locked_table.id).where(locked_table.id == locked_id).with_for_update())
+ future = pool.submit(process)
+ try:
+ future.result(timeout=5)
+ assert executor.event_buffer == {
+ TaskInstanceUuid(archiving_id): (State.SUCCESS, "terminated")
+ }
+ session.refresh(ti)
+ assert ti.state == State.RESTARTING
+ finally:
+ locker.rollback()
+
+ runner._process_executor_events(executor=executor, session=session)
+
+ assert not executor.event_buffer
+ session.refresh(ti)
+ assert (ti.working_set, ti.archived_reason) == (None, "retry")
+ successor = session.scalars(
+ select(TaskInstance).where(TaskInstance.dag_id == dr.dag_id, TaskInstance.working_set.is_(True))
+ ).one()
+ assert successor.state is None
+ executor.event_buffer[TaskInstanceUuid(archiving_id)] = State.SUCCESS, "duplicate"
+ runner._process_executor_events(executor=executor, session=session)
+ assert not executor.event_buffer
+
+ @pytest.mark.parametrize("state", [State.RUNNING, State.RESTARTING])
+ def test_heartbeat_scan_waits_for_locked_dagrun(self, dag_maker, session, state):
+ if session.bind.dialect.name == "sqlite":
+ pytest.skip("SQLite has no row locks")
+ with dag_maker(serialized=True):
+ EmptyOperator(task_id="task")
+ dr = dag_maker.create_dagrun()
+ ti = dr.get_task_instance("task", session=session)
+ runner = SchedulerJobRunner(Job(), executors=[MockExecutor(do_update=False)])
+ session.add(runner.job)
+ session.flush()
+ ti.state = state
+ ti.last_heartbeat_at = DEFAULT_DATE
+ ti.queued_by_job_id = runner.job.id
+ session.commit()
+
+ with create_session(scoped=False) as locker:
+ locker.scalar(select(DagRun.id).where(DagRun.id == dr.id).with_for_update())
+ assert runner._find_task_instances_without_heartbeats(session=session) == []
+ session.commit()
+
+ assert [row.id for row in runner._find_task_instances_without_heartbeats(session=session)] == [ti.id]
+
@pytest.mark.parametrize("event_state", [State.QUEUED, State.RUNNING])
def test_restarting_waits_for_terminal_executor_event(self, dag_maker, session, event_state):
with dag_maker(dag_id="restart_waits_for_termination"):
@@ -880,17 +952,16 @@ class TestSchedulerJob:
for record in caplog.records
)
- @pytest.mark.parametrize("event_state", [State.SUCCESS, State.UP_FOR_RETRY])
- def test_executor_event_diagnostic_distinguishes_locked_row_from_nonprocessable_state(
- self, dag_maker, session, mocker, caplog, event_state
+ def test_executor_event_with_nonprocessable_state_is_not_queried(
+ self, dag_maker, session, mocker, caplog
):
- with dag_maker(dag_id="locked_executor_event"):
+ with dag_maker(dag_id="nonprocessable_executor_event"):
EmptyOperator(task_id="task")
ti = dag_maker.create_dagrun().get_task_instance("task", session=session)
executor = MockExecutor(do_update=False)
executor._register_task(ti)
- executor.event_buffer[TaskInstanceUuid(ti.id)] = event_state, None
- scalars = mocker.patch.object(session, "scalars", autospec=True, return_value=iter(()))
+ executor.event_buffer[TaskInstanceUuid(ti.id)] = State.UP_FOR_RETRY, None
+ scalars = mocker.patch.object(session, "scalars", autospec=True)
runner = SchedulerJobRunner(Job(), executors=[executor])
runner._process_executor_events(executor=executor, session=session)
@@ -901,16 +972,54 @@ class TestSchedulerJob:
and str(ti.key) in record.message
for record in caplog.records
)
- assert any(
- "Discarding executor event" in record.message
- and "no matching task instance was returned" in record.message
- and "may be locked by another scheduler" in record.message
- for record in caplog.records
- ) == (event_state == State.SUCCESS)
- if event_state == State.SUCCESS:
- scalars.assert_called_once()
+ assert not any("Discarding executor event" in record.message for record in caplog.records)
+ scalars.assert_not_called()
+
+ @pytest.mark.parametrize("ti_fate", ["live", "missing", "archived"])
+ def test_executor_event_for_unreturned_task_instance_is_requeued_only_when_live(
+ self, dag_maker, session, mocker, caplog, ti_fate
+ ):
+ with dag_maker(dag_id="unreturned_executor_event"):
+ EmptyOperator(task_id="task")
+ ti = dag_maker.create_dagrun().get_task_instance("task", session=session)
+ ti.state = State.RUNNING
+ session.commit()
+ ti_id = ti.id
+ key = TaskInstanceUuid(ti_id)
+ executor = MockExecutor(do_update=False)
+ executor.event_buffer[key] = State.RUNNING, "external-id"
+ runner = SchedulerJobRunner(Job(), executors=[executor])
+ real_scalars = session.scalars
+
+ def skip_locked_row(*args, **kwargs):
+ if scalars.call_count > 1:
+ return real_scalars(*args, **kwargs)
+ match ti_fate:
+ case "missing":
+ session.execute(delete(TaskInstance).where(TaskInstance.id == ti_id))
+ case "archived":
+ session.execute(
+ update(TaskInstance).where(TaskInstance.id == ti_id).values(working_set=None)
+ )
+ return iter(())
+
+ scalars = mocker.patch.object(session, "scalars", autospec=True, side_effect=skip_locked_row)
+
+ returned = runner._process_executor_events(executor=executor, session=session)
+
+ discarded = any("Discarding executor event" in record.message for record in caplog.records)
+ if ti_fate == "live":
+ assert returned == 0
+ assert executor.event_buffer == {key: (State.RUNNING, "external-id")}
+ assert not discarded
+
+ runner._process_executor_events(executor=executor, session=session)
+
+ assert not executor.event_buffer
+ assert ti.external_executor_id == "external-id"
else:
- scalars.assert_not_called()
+ assert not executor.event_buffer
+ assert discarded
@mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest")
@mock.patch("airflow._shared.observability.metrics.stats._get_backend")
diff --git a/airflow-core/tests/unit/models/test_cleartasks.py b/airflow-core/tests/unit/models/test_cleartasks.py
index 5bbe5ad778b..bd12b95b880 100644
--- a/airflow-core/tests/unit/models/test_cleartasks.py
+++ b/airflow-core/tests/unit/models/test_cleartasks.py
@@ -21,11 +21,12 @@ import datetime
import random
import pytest
-from sqlalchemy import func, select, update
+from sqlalchemy import delete, func, select, update
from airflow.models.dag_version import DagVersion
from airflow.models.dagbag import DBDagBag
-from airflow.models.dagrun import DagRun
+from airflow.models.dagrun import DagRun, clear_partition_runs
+from airflow.models.dynamic_region import DynamicRegion
from airflow.models.taskinstance import TaskInstance, TaskInstance as TI, clear_task_instances
from airflow.models.taskreschedule import TaskReschedule
from airflow.providers.standard.operators.empty import EmptyOperator
@@ -38,12 +39,91 @@ from airflow.utils.types import DagRunTriggeredByType, DagRunType
from tests_common.test_utils import db
from tests_common.test_utils.dag import sync_dag_to_db
+from tests_common.test_utils.mock_operators import MockOperator
from tests_common.test_utils.taskinstance import run_task_instance
from unit.models import DEFAULT_DATE
pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
+def test_superseded_execution_is_archived_in_place_without_successor(dag_maker, session):
+ with dag_maker("superseded_clear"):
+ EmptyOperator(task_id="superseded")
+ EmptyOperator(task_id="retried")
+ dr = dag_maker.create_dagrun()
+ superseded, retried = sorted(dr.task_instances, key=lambda ti: ti.task_id, reverse=True)
+ superseded.state = retried.state = TaskInstanceState.SUCCESS
+ session.flush()
+ superseded_id, retried_id = superseded.id, retried.id
+
+ clear_task_instances([superseded, retried], session, superseded_ti_ids={superseded_id})
+
+ assert (superseded.id, superseded.working_set, superseded.archived_reason) == (
+ superseded_id,
+ None,
+ "superseded",
+ )
+ assert superseded.state == TaskInstanceState.SUCCESS
+ live = session.scalars(select(TI).where(TI.dag_id == dr.dag_id, TI.working_set.is_(True))).all()
+ assert [ti.task_id for ti in live] == ["retried"]
+ assert live[0].id != retried_id
+ assert session.get(TI, retried_id).archived_reason == "retry"
+
+
+def test_complete_restart_rejects_outcome_for_ordinary_execution(dag_maker, session):
+ with dag_maker("ordinary_restart_outcome"):
+ EmptyOperator(task_id="task")
+ dr = dag_maker.create_dagrun()
+ ti = dr.task_instances[0]
+ ti.state = TaskInstanceState.RESTARTING
+ session.flush()
+
+ with pytest.raises(ValueError, match="requires a superseded execution"):
+ ti.complete_restart(session=session, terminal_outcome=TaskInstanceState.FAILED)
+
+
+def test_partition_clear_archives_legacy_expansion_once_across_batches(dag_maker, session):
+ with dag_maker("legacy_partition_clear", serialized=True):
+ MockOperator.partial(task_id="mapped").expand(arg2=list(range(1200)))
+ dr = dag_maker.create_dagrun()
+ session.execute(delete(TI).where(TI.dag_id == dr.dag_id))
+ session.execute(delete(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id))
+ session.add_all(
+ TI(
+ dag_maker.serialized_dag.get_task("mapped"),
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_index=index,
+ state=State.SUCCESS,
+ )
+ for index in range(1200)
+ )
+ dr.partition_key = "legacy"
+ session.flush()
+
+ result = clear_partition_runs(
+ dag=None,
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ partition_key=None,
+ partition_date_start=None,
+ partition_date_end=None,
+ clear_tis=True,
+ dry_run=False,
+ session=session,
+ )
+ session.flush()
+
+ assert result == (1, 1200)
+ assert session.scalar(select(func.count()).select_from(DynamicRegion)) == 1
+ archived = session.scalars(
+ select(TI).where(TI.dag_id == dr.dag_id).execution_options(include_all_attempts=True)
+ ).all()
+ assert len(archived) == 1200
+ assert {ti.archived_reason for ti in archived} == {"superseded"}
+ assert {ti.working_set for ti in archived} == {None}
+
+
class TestClearTasks:
def test_clear_attached_attempt_keeps_pending_changes_on_archived_uuid(
self, create_task_instance, session
diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py
index f6d871a2655..8d823f29f97 100644
--- a/airflow-core/tests/unit/models/test_dag.py
+++ b/airflow-core/tests/unit/models/test_dag.py
@@ -3654,7 +3654,7 @@ class TestQueries:
for i in range(tasks_count):
EmptyOperator(task_id=f"dummy_task_{i}", owner="test", dag=dag)
scheduler_dag = sync_dag_to_db(dag)
- with assert_queries_count(5):
+ with assert_queries_count(6):
scheduler_dag.create_dagrun(
run_id="test_dagrun_query_count",
run_type=DagRunType.MANUAL,
diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py
index 3c33a19bbfd..3d4f968efd6 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -5609,7 +5609,7 @@ class TestClearPartitionRuns:
with (
mock.patch("airflow.models.dagrun._TI_CHUNK_SIZE", 6),
- mock.patch("airflow.models.dagrun.clear_task_instances", autospec=True) as mock_cti,
+ mock.patch("airflow.models.loop_clear.clear_task_instances_for_runs", autospec=True) as mock_cti,
):
cleared, tis = clear_partition_runs(
dag=serialized_dag,
@@ -5655,7 +5655,7 @@ class TestClearPartitionRuns:
with (
mock.patch("airflow.models.dagrun._TI_CHUNK_SIZE", 6),
- mock.patch("airflow.models.dagrun.clear_task_instances", autospec=True) as mock_cti,
+ mock.patch("airflow.models.loop_clear.clear_task_instances_for_runs", autospec=True) as mock_cti,
):
clear_partition_runs(
dag=serialized_dag,
diff --git a/airflow-core/tests/unit/models/test_dynamic_region.py b/airflow-core/tests/unit/models/test_dynamic_region.py
index df32cc23538..8e1c7e741a2 100644
--- a/airflow-core/tests/unit/models/test_dynamic_region.py
+++ b/airflow-core/tests/unit/models/test_dynamic_region.py
@@ -406,7 +406,11 @@ def _ti_search_plans(session, statements):
plans = []
for statement in statements:
rows = session.execute(text(f"EXPLAIN QUERY PLAN {statement}")).all()
- plans.extend(row[-1] for row in rows if "task_instance" in row[-1] and "SEARCH" in row[-1])
+ plans.extend(
+ row[-1]
+ for row in rows
+ if "task_instance" in row[-1] and ("SEARCH" in row[-1] or "SCAN" in row[-1])
+ )
return plans
@@ -434,3 +438,9 @@ def test_public_lookups_use_the_task_instance_unique_key(dag_maker, session):
plans = _ti_search_plans(session, statements)
assert any("region_id=? AND region_index=?" in plan for plan in plans), (name, plans)
assert not any("ANY(" in plan for plan in plans), (name, plans)
+
+ with capture_orm_selects("task_instance") as statements:
+ dr._reconcile_legacy_expansions(session=session)
+ plans = _ti_search_plans(session, statements)
+ assert plans, statements
+ assert all("(dag_id=?" in plan for plan in plans), (statements, plans)
diff --git a/airflow-core/tests/unit/models/test_loop_clear.py b/airflow-core/tests/unit/models/test_loop_clear.py
new file mode 100644
index 00000000000..2a4c42431e7
--- /dev/null
+++ b/airflow-core/tests/unit/models/test_loop_clear.py
@@ -0,0 +1,947 @@
+# 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 concurrent.futures import ThreadPoolExecutor
+from threading import Event
+
+import pytest
+from sqlalchemy import delete, event, select, update
+from sqlalchemy.orm import Session
+from sqlalchemy.sql import Select
+
+from airflow._shared.timezones import timezone
+from airflow.api_fastapi.core_api.datamodels.task_instances import PatchTaskInstanceBody
+from airflow.api_fastapi.core_api.services.public.dag_run import perform_clear_dag_run
+from airflow.api_fastapi.core_api.services.public.task_instances import _patch_selected_task_state
+from airflow.cli.commands.dag_command import _bulk_clear_runs
+from airflow.exceptions import AirflowClearRunningTaskException
+from airflow.models.dag_version import DagVersion
+from airflow.models.dagbag import DBDagBag
+from airflow.models.dagrun import DagRun, clear_partition_runs
+from airflow.models.dynamic_region import LOOP_DECISION_KEY, DynamicRegion
+from airflow.models.loop_clear import (
+ LoopClearScope,
+ apply_loop_clear_scope,
+ clear_loop_task_instances,
+ loop_gate_waits_for_archival,
+ select_loop_clear_scope,
+)
+from airflow.models.serialized_dag import SerializedDagModel
+from airflow.models.task_coordinates import TaskCoordinateResolver
+from airflow.models.taskinstance import TaskInstance, clear_task_instances
+from airflow.models.xcom import XComModelV2
+from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.providers.standard.operators.python import PythonOperator
+from airflow.sdk import task_group
+from airflow.sdk.definitions._internal.loop import create_loop
+from airflow.utils.session import create_session
+from airflow.utils.state import State
+
+pytestmark = [pytest.mark.db_test, pytest.mark.need_serialized_dag]
+
+
+@pytest.fixture
+def loop_run(dag_maker, session):
+ @task_group
+ def body():
+ prepare = EmptyOperator(task_id="prepare")
+ process = PythonOperator.partial(task_id="process", python_callable=list).expand(op_kwargs=[{}, {}])
+ consume = EmptyOperator(task_id="consume")
+ prepare >> process >> consume
+
+ with dag_maker(serialized=True) as dag:
+ loop = create_loop(body, max_iterations=5)
+ loop >> EmptyOperator(task_id="outside")
+ dr = dag_maker.create_dagrun()
+ root = session.scalar(select(DynamicRegion).where(DynamicRegion.node_id == loop.group_id))
+ for index in range(1, 5):
+ child = DynamicRegion(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id="body.process",
+ parent_region_id=root.id,
+ parent_region_index=index,
+ )
+ session.add(child)
+ session.flush()
+ for task in loop.iter_tasks():
+ mapped = task.task_id == "body.process"
+ for slot in range(2) if mapped else [index]:
+ session.add(
+ TaskInstance(
+ task=task,
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=child.id if mapped else root.id,
+ region_index=slot,
+ )
+ )
+ session.flush()
+ tis = list(dr.get_task_instances(session=session))
+ regions = {r.id: r for r in session.scalars(select(DynamicRegion))}
+
+ def iteration(ti):
+ return regions[ti.region_id].parent_region_index if ti.task_id == "body.process" else ti.region_index
+
+ return dr, dag, loop, root, tis, iteration
+
+
+@pytest.mark.parametrize("downstream", [False, True])
+@pytest.mark.parametrize("later", [False, True])
+def test_loop_clear_controls_select_exact_live_executions(loop_run, session, downstream, later):
+ dr, dag, loop, root, tis, iteration = loop_run
+ selected = next(
+ ti for ti in tis if ti.task_id == "body.process" and iteration(ti) == 2 and ti.region_index == 0
+ )
+ before = [(ti.id, ti.region_id, ti.region_index, ti.state, ti.working_set) for ti in tis]
+
+ scope = select_loop_clear_scope(
+ [selected], downstream=downstream, later_loop_iterations=later, session=session
+ )
+
+ expected_retry = {selected.id}
+ if downstream:
+ expected_retry.update(
+ ti.id
+ for ti in tis
+ if ti.task_id == "outside"
+ or (iteration(ti) == 2 and ti.task_id in {"body.consume", loop.gate_task_id})
+ )
+ assert scope.retry_ids == expected_retry
+ assert scope.archive_ids == {
+ ti.id for ti in tis if ti.task_id != "outside" and iteration(ti) > 2 and downstream and later
+ }
+ assert [(ti.id, ti.region_id, ti.region_index, ti.state, ti.working_set) for ti in tis] == before
+ assert not session.new
+ assert not session.deleted
+ assert not session.dirty
+
+
+@pytest.mark.parametrize("whole", [False, True])
+def test_whole_mapped_task_selection_stays_in_selected_iteration(loop_run, session, whole):
+ dr, dag, loop, root, tis, iteration = loop_run
+ selected = next(
+ ti for ti in tis if ti.task_id == "body.process" and iteration(ti) == 2 and ti.region_index == 0
+ )
+
+ scope = select_loop_clear_scope(
+ [selected, selected],
+ whole_expansion_ids={selected.id} if whole else (),
+ downstream=False,
+ session=session,
+ )
+
+ assert scope.retry_ids == {
+ ti.id
+ for ti in tis
+ if ti.task_id == selected.task_id and iteration(ti) == 2 and (whole or ti.region_index == 0)
+ }
+ assert not scope.archive_ids
+
+
+def test_direct_gate_clear_archives_later_passes_without_downstream(loop_run, session):
+ dr, dag, loop, root, tis, iteration = loop_run
+ gate = next(ti for ti in tis if ti.task_id == loop.gate_task_id and ti.region_index == 2)
+
+ scope = select_loop_clear_scope([gate], downstream=False, session=session)
+
+ assert scope.retry_ids == {gate.id}
+ assert scope.archive_ids == {ti.id for ti in tis if ti.task_id != "outside" and iteration(ti) > 2}
+
+
+def test_loop_clear_uses_execution_pinned_graph_after_latest_definition_changes(loop_run, dag_maker, session):
+ dr, dag, loop, root, tis, iteration = loop_run
+ selected = next(ti for ti in tis if ti.task_id == "body.consume" and ti.region_index == 2)
+ with dag_maker(serialized=True, session=session):
+ EmptyOperator(task_id="replacement")
+
+ scope = select_loop_clear_scope([selected], later_loop_iterations=False, session=session)
+
+ assert scope.retry_ids == {
+ ti.id
+ for ti in tis
+ if ti.task_id == "outside"
+ or (iteration(ti) == 2 and ti.task_id in {selected.task_id, loop.gate_task_id})
+ }
+ assert not scope.archive_ids
+
+
+@pytest.mark.parametrize("width", [1, 3])
+@pytest.mark.parametrize("whole", [False, True])
+def test_loop_clear_preserves_mapped_group_relevant_indexes(dag_maker, session, width, whole):
+ @task_group
+ def mapped(value):
+ PythonOperator(task_id="first", python_callable=list, op_kwargs={"value": value}) >> EmptyOperator(
+ task_id="last"
+ )
+
+ @task_group
+ def body():
+ mapped.expand(value=list(range(width)))
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=3)
+ dr = dag_maker.create_dagrun()
+ tis = dr.get_task_instances(session=session)
+ selected = next(ti for ti in tis if ti.task_id == "body.mapped.first" and ti.region_index == 0)
+
+ scope = select_loop_clear_scope(
+ [selected], whole_expansion_ids={selected.id} if whole else (), session=session
+ )
+
+ assert scope.retry_ids == {
+ ti.id for ti in tis if ti.task_id == loop.gate_task_id or whole or ti.region_index == 0
+ }
+ assert not scope.archive_ids
+
+
+def test_loop_clear_does_not_resurrect_archived_passes_after_repeated_selection(loop_run, session):
+ dr, dag, loop, root, tis, iteration = loop_run
+ selected = next(ti for ti in tis if ti.task_id == "body.process" and iteration(ti) == 2)
+ first = select_loop_clear_scope([selected], whole_expansion_ids={selected.id}, session=session)
+ replacement = DynamicRegion(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id=loop.group_id,
+ forked_from_region_id=root.id,
+ resumes_from_index=3,
+ )
+ session.add(replacement)
+ session.flush()
+ for ti in tis:
+ if ti.id in first.archive_ids:
+ ti.archive(reason="superseded", session=session)
+ session.flush()
+ replacement_gate = TaskInstance(
+ task=dag.get_task(loop.gate_task_id),
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=replacement.id,
+ region_index=3,
+ )
+ session.add(replacement_gate)
+ session.flush()
+ consume = next(ti for ti in tis if ti.task_id == "body.consume" and ti.region_index == 2)
+
+ second = select_loop_clear_scope([consume], later_loop_iterations=False, session=session)
+ rewind = select_loop_clear_scope([consume], session=session)
+
+ assert second.retry_ids == {
+ ti.id
+ for ti in tis
+ if ti.task_id == "outside"
+ or (iteration(ti) == 2 and ti.task_id in {consume.task_id, loop.gate_task_id})
+ }
+ assert not second.archive_ids
+ assert rewind.archive_ids == {replacement_gate.id}
+ assert not (first.archive_ids & (second.retry_ids | rewind.archive_ids))
+
+
+def test_loop_clear_rejects_deleted_selected_execution(loop_run, session):
+ dr, dag, loop, root, tis, iteration = loop_run
+ selected = next(ti for ti in tis if ti.task_id == "body.consume")
+ selected.archive(reason="superseded", session=session)
+
+ with pytest.raises(ValueError, match="no longer live"):
+ select_loop_clear_scope([selected], session=session)
+
+
+def test_ordinary_seed_without_relatives_selects_only_itself(loop_run, session):
+ dr, dag, loop, root, tis, iteration = loop_run
+ selected = next(ti for ti in tis if ti.task_id == "outside")
+
+ scope = select_loop_clear_scope([selected], downstream=False, session=session)
+ assert scope.retry_ids == {selected.id}
+ assert not scope.archive_ids
+
+
+def test_mixed_clear_scope_applies_together_and_empty_scope_is_noop(loop_run, session):
+ dr, dag, loop, root, tis, iteration = loop_run
+ outside = next(ti for ti in tis if ti.task_id == "outside")
+ gate = next(ti for ti in tis if ti.task_id == loop.gate_task_id and ti.region_index == 2)
+ scope = select_loop_clear_scope([outside, gate], downstream=False, session=session)
+ outside.state = gate.state = State.SUCCESS
+ session.flush()
+
+ cleared = apply_loop_clear_scope(scope, session=session, dag_run_state=False)
+ assert apply_loop_clear_scope(LoopClearScope(frozenset(), frozenset()), session=session) == []
+
+ assert {ti.id for ti in cleared if ti.working_set is None} == scope.archive_ids
+ assert {ti.id for ti in cleared if ti.working_set is True} == {
+ ti.id for ti in dr.get_task_instances(session=session) if ti.id not in {t.id for t in tis}
+ }
+ assert {ti.archived_reason for ti in cleared if ti.working_set is None} == {"superseded"}
+ assert outside.working_set is None
+ assert gate.working_set is None
+ assert {ti.id for ti in dr.get_task_instances(session=session)} - {t.id for t in tis} == {
+ ti.id for ti in cleared if ti.working_set is True
+ }
+
+
+def test_running_clear_rejection_does_not_allocate_fork(loop_run, session):
+ dr, dag, loop, root, tis, iteration = loop_run
+ gate = next(ti for ti in tis if ti.task_id == loop.gate_task_id and ti.region_index == 2)
+ gate.state = State.RUNNING
+ session.flush()
+ scope = select_loop_clear_scope([gate], downstream=False, session=session)
+
+ with pytest.raises(AirflowClearRunningTaskException):
+ apply_loop_clear_scope(scope, session=session, prevent_running_task=True)
+
+ assert not session.scalar(
+ select(DynamicRegion.id).where(DynamicRegion.forked_from_region_id.is_not(None))
+ )
+
+
+@pytest.mark.parametrize("downstream", [False, True])
+def test_loop_upstream_scope_stays_in_iteration_before_downstream_expansion(loop_run, session, downstream):
+ dr, dag, loop, root, tis, iteration = loop_run
+ consume = next(ti for ti in tis if ti.task_id == "body.consume" and ti.region_index == 2)
+
+ scope = select_loop_clear_scope([consume], upstream=True, downstream=downstream, session=session)
+
+ assert scope.retry_ids == {
+ ti.id
+ for ti in tis
+ if (
+ ti.task_id != "outside" and iteration(ti) == 2 and (downstream or ti.task_id != loop.gate_task_id)
+ )
+ or (downstream and ti.task_id == "outside")
+ }
+ assert scope.archive_ids == {
+ ti.id for ti in tis if downstream and ti.task_id != "outside" and iteration(ti) > 2
+ }
+
+
+def test_outside_seed_upstream_selects_all_current_loop_occurrences(loop_run, session):
+ dr, dag, loop, root, tis, iteration = loop_run
+ outside = next(ti for ti in tis if ti.task_id == "outside")
+
+ scope = select_loop_clear_scope([outside], upstream=True, downstream=False, session=session)
+
+ assert scope.retry_ids == {ti.id for ti in tis if ti.task_id == "outside" or iteration(ti) == 0}
+ assert scope.archive_ids == {ti.id for ti in tis if ti.task_id != "outside" and iteration(ti) > 0}
+
+
+@pytest.mark.parametrize("upstream", [False, True])
+def test_outside_predecessor_selects_all_mapped_loop_occurrences(dag_maker, session, upstream):
+ @task_group
+ def body():
+ PythonOperator.partial(task_id="process", python_callable=list).expand(op_kwargs=[{}, {}])
+
+ with dag_maker(serialized=True):
+ before = EmptyOperator(task_id="before")
+ loop = create_loop(body, max_iterations=3)
+ before >> loop
+ dr = dag_maker.create_dagrun()
+ root = session.scalar(select(DynamicRegion).where(DynamicRegion.node_id == loop.group_id))
+ child = DynamicRegion(
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ node_id="body.process",
+ parent_region_id=root.id,
+ parent_region_index=1,
+ )
+ session.add(child)
+ session.flush()
+ for task in loop.iter_tasks():
+ mapped = task.task_id == "body.process"
+ for index in range(2) if mapped else [1]:
+ session.add(
+ TaskInstance(
+ task=task,
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=child.id if mapped else root.id,
+ region_index=index,
+ )
+ )
+ session.flush()
+ tis = dr.get_task_instances(session=session)
+ seed = next(
+ ti
+ for ti in tis
+ if ti.task_id == ("body.process" if upstream else "before")
+ and ti.region_index == (0 if upstream else -1)
+ )
+
+ scope = select_loop_clear_scope(
+ [seed],
+ upstream=upstream,
+ downstream=True,
+ later_loop_iterations=False,
+ session=session,
+ )
+
+ assert scope.retry_ids == {ti.id for ti in tis}
+ assert not scope.archive_ids
+
+
+def test_loop_clear_refreshes_pending_execution_version_after_concurrent_clear(loop_run, dag_maker, session):
+ dr, dag, loop, root, tis, iteration = loop_run
+ selected = next(ti for ti in tis if ti.task_id == "body.consume" and ti.region_index == 2)
+ old_version = selected.dag_version_id
+ selected_id = selected.id
+
+ @task_group
+ def body():
+ prepare = EmptyOperator(task_id="prepare")
+ process = PythonOperator.partial(task_id="process", python_callable=list).expand(op_kwargs=[{}, {}])
+ consume = EmptyOperator(task_id="consume")
+ prepare >> process >> consume >> EmptyOperator(task_id="new")
+
+ with dag_maker(serialized=True, session=session):
+ create_loop(body, max_iterations=5) >> EmptyOperator(task_id="outside")
+ new_version = DagVersion.get_latest_version(dr.dag_id, session=session).id
+ session.commit()
+ assert selected.dag_version_id == old_version
+ with create_session(scoped=False) as other_session:
+ pending = other_session.get(TaskInstance, selected_id)
+ clear_task_instances([pending], other_session, run_on_latest_version=True)
+ assert pending.id == selected_id
+ assert pending.dag_version_id == new_version
+ pending.dag_run.verify_integrity(session=other_session, dag_version_id=new_version)
+ new_id = other_session.scalar(
+ select(TaskInstance.id).where(
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.run_id == dr.run_id,
+ TaskInstance.task_id == "body.new",
+ TaskInstance.region_index == 2,
+ )
+ )
+ assert new_id is not None
+ assert selected.dag_version_id == old_version
+
+ scope = select_loop_clear_scope([selected], later_loop_iterations=False, session=session)
+
+ assert scope.retry_ids == {
+ ti.id
+ for ti in tis
+ if ti.task_id == "outside"
+ or (iteration(ti) == 2 and ti.task_id in {selected.task_id, loop.gate_task_id})
+ } | {new_id}
+ assert selected.dag_version_id == new_version
+
+
+def test_loop_clear_mixes_whole_and_single_index_for_same_task_in_different_passes(loop_run, session):
+ dr, dag, loop, root, tis, iteration = loop_run
+ whole = next(
+ ti for ti in tis if ti.task_id == "body.process" and iteration(ti) == 1 and ti.region_index == 0
+ )
+ single = next(
+ ti for ti in tis if ti.task_id == "body.process" and iteration(ti) == 2 and ti.region_index == 0
+ )
+
+ scope = select_loop_clear_scope(
+ [whole, single], whole_expansion_ids={whole.id}, downstream=False, session=session
+ )
+
+ assert scope.retry_ids == {ti.id for ti in tis if ti.task_id == whole.task_id and iteration(ti) == 1} | {
+ single.id
+ }
+ assert not scope.archive_ids
+
+
+@pytest.fixture
+def completed_loop(dag_maker, session):
+ @task_group
+ def body():
+ (
+ EmptyOperator(task_id="prepare")
+ >> EmptyOperator(task_id="process")
+ >> EmptyOperator(task_id="consume")
+ )
+
+ with dag_maker(serialized=True):
+ loop = create_loop(body, max_iterations=5, until=lambda loop: True)
+ dr = dag_maker.create_dagrun()
+ session.scalar(select(DagRun).where(DagRun.id == dr.id).with_for_update())
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+
+ def complete_pass(index, decision):
+ members = [ti for ti in dr.get_task_instances(session=session) if ti.region_index == index]
+ for ti in members:
+ if ti.state != State.SUCCESS:
+ ti.state = State.SUCCESS
+ ti.start_date = ti.end_date = timezone.utcnow()
+ gate = next(ti for ti in members if ti.task_id == loop.gate_task_id)
+ session.add(XComModelV2(task_instance_id=gate.id, key=LOOP_DECISION_KEY, value=decision))
+ session.flush()
+ group, _ = resolver.loop_context(gate)
+ dr.complete_loop_gate(gate, group, State.SUCCESS, session=session)
+ session.flush()
+
+ for index in range(5):
+ complete_pass(index, "stop" if index == 4 else "continue")
+ return dr, loop, complete_pass
+
+
+def test_integrity_does_not_create_members_at_draining_gate_coordinates(completed_loop, dag_maker, session):
+ dr, loop, complete_pass = completed_loop
+ gates = {
+ ti.region_index: ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == loop.gate_task_id
+ }
+ gates[4].state = State.RUNNING
+ session.flush()
+ clear_loop_task_instances([gates[2]], downstream=False, session=session)
+
+ @task_group
+ def body():
+ (
+ EmptyOperator(task_id="prepare")
+ >> EmptyOperator(task_id="process")
+ >> EmptyOperator(task_id="consume")
+ >> EmptyOperator(task_id="new")
+ )
+
+ with dag_maker(serialized=True, session=session):
+ create_loop(body, max_iterations=5, until=lambda loop: True)
+ version = DagVersion.get_latest_version(dr.dag_id, session=session).id
+ dr.dag = DBDagBag().get_dag(version_id=version, session=session)
+
+ dr.verify_integrity(session=session, dag_version_id=version)
+
+ assert gates[4].state == State.RESTARTING
+ assert set(
+ session.scalars(
+ select(TaskInstance.region_index).where(
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.run_id == dr.run_id,
+ TaskInstance.task_id == "body.new",
+ )
+ )
+ ) == {0, 1, 2}
+
+
+@pytest.mark.parametrize("decision", ["stop", "continue"])
+def test_repeated_selective_loop_clear_preserves_unselected_work(completed_loop, session, decision):
+ dr, loop, complete_pass = completed_loop
+ original = {(ti.task_id, ti.region_index): ti.id for ti in dr.get_task_instances(session=session)}
+ process = session.get(TaskInstance, original["body.process", 2])
+
+ clear_loop_task_instances([process], session=session)
+ complete_pass(2, "continue")
+ complete_pass(3, "stop")
+ first = {(ti.task_id, ti.region_index): ti.id for ti in dr.get_task_instances(session=session)}
+ consume = session.get(TaskInstance, first["body.consume", 2])
+ clear_loop_task_instances([consume], later_loop_iterations=False, session=session)
+ complete_pass(2, decision)
+
+ current = {(ti.task_id, ti.region_index): ti.id for ti in dr.get_task_instances(session=session)}
+ assert current["body.prepare", 2] == original["body.prepare", 2]
+ assert current["body.process", 2] == first["body.process", 2] != original["body.process", 2]
+ assert current["body.consume", 2] != first["body.consume", 2]
+ assert current[loop.gate_task_id, 2] != first[loop.gate_task_id, 2]
+ assert {key: value for key, value in current.items() if key[1] == 3} == {
+ key: value for key, value in first.items() if key[1] == 3
+ }
+ assert not any(index == 4 for _, index in current)
+ assert session.get(TaskInstance, original["body.prepare", 2]).working_set is True
+ original_four = session.get(TaskInstance, original[loop.gate_task_id, 4])
+ assert (original_four.working_set, original_four.archived_reason) == (None, "superseded")
+ assert session.get(TaskInstance, original["body.process", 2]).archived_reason == "retry"
+
+
+def test_gate_rerun_after_chained_forks_keeps_later_passes_in_later_regions(completed_loop, session):
+ dr, loop, complete_pass = completed_loop
+
+ def gate_at(index):
+ return next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == loop.gate_task_id and ti.region_index == index
+ )
+
+ clear_loop_task_instances([gate_at(2)], downstream=False, session=session)
+ complete_pass(2, "continue")
+ complete_pass(3, "continue")
+ clear_loop_task_instances([gate_at(1)], downstream=False, session=session)
+ complete_pass(1, "continue")
+ complete_pass(2, "continue")
+ before = {
+ (ti.task_id, ti.region_index): (ti.id, ti.region_id) for ti in dr.get_task_instances(session=session)
+ }
+ clear_loop_task_instances([gate_at(1)], downstream=False, later_loop_iterations=False, session=session)
+
+ complete_pass(1, "continue")
+
+ current = {
+ (ti.task_id, ti.region_index): (ti.id, ti.region_id) for ti in dr.get_task_instances(session=session)
+ }
+ assert {key: value for key, value in current.items() if key[1] >= 2} == {
+ key: value for key, value in before.items() if key[1] >= 2
+ }
+
+
+def clear_run_through_dag_clear(dr, session):
+ dr.dag.clear(run_id=dr.run_id, session=session)
+
+
+def clear_run_through_api_service(dr, session):
+ perform_clear_dag_run(
+ session=session,
+ dag=dr.dag,
+ dag_run=dr,
+ dag_id=dr.dag_id,
+ only_failed=False,
+ only_new=False,
+ run_on_latest_version=False,
+ note=None,
+ user=None,
+ )
+
+
+def clear_run_through_partition_clear(dr, session):
+ clear_partition_runs(
+ dag=dr.dag,
+ dag_id=dr.dag_id,
+ run_id=dr.run_id,
+ partition_key=None,
+ partition_date_start=None,
+ partition_date_end=None,
+ clear_tis=True,
+ dry_run=False,
+ session=session,
+ )
+
+
+def clear_run_through_cli_bulk_clear(dr, session):
+ _bulk_clear_runs(dr.dag_id, [dr.run_id], only_failed=False, only_running=False, session=session)
+
+
+@pytest.mark.parametrize(
+ "clear_run",
+ [
+ clear_run_through_dag_clear,
+ clear_run_through_api_service,
+ clear_run_through_partition_clear,
+ clear_run_through_cli_bulk_clear,
+ ],
+)
+class TestWholeRunClear:
+ def test_archives_every_later_pass_and_regenerates_from_the_first(
+ self, completed_loop, session, clear_run
+ ):
+ dr, loop, complete_pass = completed_loop
+ original = {(ti.task_id, ti.region_index): ti.id for ti in dr.get_task_instances(session=session)}
+
+ clear_run(dr, session)
+
+ live = dr.get_task_instances(session=session)
+ assert {ti.region_index for ti in live} == {0}
+ for (task_id, index), ti_id in original.items():
+ archived = session.get(TaskInstance, ti_id)
+ if index == 0:
+ assert (archived.working_set, archived.archived_reason) == (None, "retry"), task_id
+ else:
+ assert (archived.working_set, archived.archived_reason) == (None, "superseded"), task_id
+ fork = session.scalar(select(DynamicRegion).where(DynamicRegion.forked_from_region_id.is_not(None)))
+ assert fork.resumes_from_index == 1
+
+ complete_pass(0, "continue")
+
+ regenerated = [ti for ti in dr.get_task_instances(session=session) if ti.region_index == 1]
+ assert len(regenerated) == 4
+ assert {ti.region_id for ti in regenerated} == {fork.id}
+
+ def test_rerun_can_shorten_the_loop(self, completed_loop, session, clear_run):
+ dr, loop, complete_pass = completed_loop
+
+ clear_run(dr, session)
+ complete_pass(0, "stop")
+
+ assert {ti.region_index for ti in dr.get_task_instances(session=session)} == {0}
+
+
+def test_selected_state_patch_locks_dag_run_before_task_instances(completed_loop, session):
+ dr, loop, _ = completed_loop
+ selected = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == "body.consume" and ti.region_index == 2
+ )
+ locked: list[str] = []
+
+ @event.listens_for(session, "do_orm_execute")
+ def record_locks(orm_execute_state):
+ statement = orm_execute_state.statement
+ if isinstance(statement, Select) and statement._for_update_arg is not None:
+ locked.append(str(statement))
+
+ try:
+ _patch_selected_task_state(
+ [selected],
+ PatchTaskInstanceBody(new_state="failed"),
+ {"new_state": State.FAILED},
+ session=session,
+ commit=True,
+ )
+ finally:
+ event.remove(session, "do_orm_execute", record_locks)
+
+ tables = ["dag_run" if "FROM dag_run" in sql else "task_instance" for sql in locked]
+ assert "task_instance" in tables
+ assert tables[0] == "dag_run"
+
+
+@pytest.mark.backend("mysql", "postgres")
+def test_selected_state_patch_and_clear_take_locks_in_the_same_order(completed_loop, session):
+ dr, loop, _ = completed_loop
+ selected = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == "body.consume" and ti.region_index == 2
+ )
+ selected_id = selected.id
+ next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == loop.gate_task_id and ti.region_index == 2
+ ).state = State.UPSTREAM_FAILED
+ session.commit()
+ bind = session.get_bind()
+ clear_holds_run_lock = Event()
+
+ def patch_state():
+ with Session(bind=bind) as patch_session:
+ locks = 0
+ waited = False
+
+ @event.listens_for(patch_session, "do_orm_execute")
+ def pause_after_first_lock(orm_execute_state):
+ nonlocal locks, waited
+ statement = orm_execute_state.statement
+ if isinstance(statement, Select) and statement._for_update_arg is not None:
+ locks += 1
+ elif locks and not waited:
+ waited = True
+ clear_holds_run_lock.wait(timeout=3)
+
+ _patch_selected_task_state(
+ [patch_session.get(TaskInstance, selected_id)],
+ PatchTaskInstanceBody(new_state="failed"),
+ {"new_state": State.FAILED},
+ session=patch_session,
+ commit=True,
+ )
+ patch_session.commit()
+
+ def clear():
+ with Session(bind=bind) as clear_session:
+ run_locked = False
+
+ @event.listens_for(clear_session, "do_orm_execute")
+ def signal_after_run_lock(orm_execute_state):
+ nonlocal run_locked
+ statement = orm_execute_state.statement
+ if run_locked:
+ clear_holds_run_lock.set()
+ elif isinstance(statement, Select) and statement._for_update_arg is not None:
+ run_locked = True
+
+ clear_task_instances([clear_session.get(TaskInstance, selected_id)], session=clear_session)
+ clear_session.commit()
+
+ with ThreadPoolExecutor(max_workers=2) as pool:
+ futures = [pool.submit(patch_state), pool.submit(clear)]
+ for future in futures:
+ future.result(timeout=30)
+
+
+def test_archive_decides_from_the_state_committed_after_the_lock(completed_loop, session):
+ dr, loop, _ = completed_loop
+ stale = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == "body.consume" and ti.region_index == 2
+ )
+ stale.state, stale.start_date, stale.end_date = State.QUEUED, None, None
+ session.flush()
+ session.execute(
+ update(TaskInstance)
+ .where(TaskInstance.id == stale.id)
+ .values(state=State.RUNNING, start_date=timezone.utcnow())
+ .execution_options(synchronize_session=False)
+ )
+
+ stale.archive(reason="superseded", session=session)
+
+ session.refresh(stale)
+ assert (stale.working_set, stale.state) == (None, State.FAILED)
+ assert stale.start_date is not None
+ assert stale.end_date is not None
+
+
+def test_archival_wait_needs_no_dag_definition_before_any_fork(completed_loop, session, mocker):
+ dr, loop, _ = completed_loop
+ gate = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == loop.gate_task_id and ti.region_index == 2
+ )
+ dag_bag = mocker.patch("airflow.models.loop_clear.DBDagBag", autospec=True)
+ statements = []
+
+ @event.listens_for(session, "do_orm_execute")
+ def count_statements(orm_execute_state):
+ statements.append(orm_execute_state.statement)
+
+ try:
+ assert loop_gate_waits_for_archival(gate, session=session) is False
+ finally:
+ event.remove(session, "do_orm_execute", count_statements)
+
+ assert len(statements) == 1
+ dag_bag.assert_not_called()
+
+
+def test_archival_wait_loads_no_task_instance_entities(completed_loop, session):
+ dr, loop, _ = completed_loop
+ gates = {
+ ti.region_index: ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == loop.gate_task_id
+ }
+ gates[4].state = State.RUNNING
+ session.flush()
+ clear_loop_task_instances([gates[2]], downstream=False, session=session)
+ gate_id = gates[2].id
+ session.expire_all()
+ session.expunge_all()
+ loaded = []
+
+ @event.listens_for(TaskInstance, "load")
+ def count_loaded(target, context):
+ loaded.append(target.id)
+
+ try:
+ assert loop_gate_waits_for_archival(session.get(TaskInstance, gate_id), session=session) is True
+ finally:
+ event.remove(TaskInstance, "load", count_loaded)
+
+ assert len(loaded) == 1
+
+
+def test_archival_wait_ignores_terminating_passes_at_or_before_the_gate(completed_loop, session):
+ dr, loop, _ = completed_loop
+ members = {(ti.task_id, ti.region_index): ti for ti in dr.get_task_instances(session=session)}
+ members[loop.gate_task_id, 2].state = State.RUNNING
+ session.flush()
+ clear_loop_task_instances([members[loop.gate_task_id, 0]], downstream=False, session=session)
+ fork = session.scalar(select(DynamicRegion).where(DynamicRegion.forked_from_region_id.is_not(None)))
+ gate_task = next(task for task in loop.iter_tasks() if task.task_id == loop.gate_task_id)
+ new_gates = {}
+ for index in (1, 3):
+ new_gates[index] = TaskInstance(
+ task=gate_task,
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_id=fork.id,
+ region_index=index,
+ state=State.RUNNING,
+ )
+ session.add(new_gates[index])
+ session.flush()
+
+ assert members[loop.gate_task_id, 2].state == State.RESTARTING
+ assert loop_gate_waits_for_archival(new_gates[1], session=session) is True
+ assert loop_gate_waits_for_archival(new_gates[3], session=session) is False
+
+
+def test_repeated_rewind_appends_empty_forks_and_generates_in_latest_region(completed_loop, session):
+ dr, loop, complete_pass = completed_loop
+ original_region = next(ti.region_id for ti in dr.get_task_instances(session=session))
+ for index in (2, 1):
+ gate = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == loop.gate_task_id and ti.region_index == index
+ )
+ clear_loop_task_instances([gate], downstream=False, session=session)
+ complete_pass(index, "stop")
+ regions = list(session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)))
+ first = next(region for region in regions if region.forked_from_region_id == original_region)
+ second = next(region for region in regions if region.forked_from_region_id == first.id)
+ assert (first.resumes_from_index, second.resumes_from_index) == (3, 2)
+ assert not session.scalar(
+ select(TaskInstance.id).where(TaskInstance.region_id.in_([first.id, second.id]))
+ )
+ gate = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == loop.gate_task_id and ti.region_index == 1
+ )
+
+ clear_loop_task_instances([gate], downstream=False, session=session)
+ complete_pass(1, "continue")
+
+ third = session.scalar(select(DynamicRegion).where(DynamicRegion.forked_from_region_id == second.id))
+ generated = [ti for ti in dr.get_task_instances(session=session) if ti.region_index == 2]
+ assert len(generated) == 4
+ assert {ti.region_id for ti in generated} == {third.id}
+
+
+def test_ordinary_loop_task_clear_does_not_allocate_fork(completed_loop, session):
+ dr, loop, complete_pass = completed_loop
+ selected = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == "body.consume" and ti.region_index == 2
+ )
+ coordinate = selected.region_id, selected.region_index
+ original_id = selected.id
+
+ clear_loop_task_instances([selected], downstream=False, session=session)
+
+ successor = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == dr.dag_id,
+ TaskInstance.task_id == "body.consume",
+ TaskInstance.region_index == 2,
+ TaskInstance.working_set.is_(True),
+ )
+ ).one()
+ assert successor.id != original_id
+ assert (successor.region_id, successor.region_index) == coordinate
+ assert (selected.working_set, selected.archived_reason) == (None, "retry")
+ assert len(session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)).all()) == 1
+
+
+def test_mapped_archival_ack_uses_fork_facts_when_definition_is_missing(loop_run, session):
+ dr, dag, loop, root, tis, iteration = loop_run
+ archiving = next(
+ ti for ti in tis if ti.task_id == "body.process" and iteration(ti) == 3 and ti.region_index == 0
+ )
+ archiving.state = State.RUNNING
+ session.flush()
+ gate = next(ti for ti in tis if ti.task_id == loop.gate_task_id and ti.region_index == 2)
+ clear_loop_task_instances([gate], downstream=False, session=session)
+ session.execute(
+ delete(SerializedDagModel).where(SerializedDagModel.dag_version_id == archiving.dag_version_id)
+ )
+ session.flush()
+ session.expire_all()
+ assert DBDagBag().get_dag(archiving.dag_version_id, session=session) is None
+ archiving_id = archiving.id
+
+ archiving.complete_restart(session=session)
+
+ archived = session.get(TaskInstance, archiving_id)
+ assert (archived.working_set, archived.archived_reason) == (None, "superseded")
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py
index f582ddc3a10..ce3bbcaf577 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -62,6 +62,7 @@ from airflow.models.dag_version import DagVersion
from airflow.models.dagrun import DagRun
from airflow.models.deadline import Deadline
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
+from airflow.models.dynamic_region import DynamicRegion
from airflow.models.pool import Pool
from airflow.models.renderedtifields import RenderedTaskInstanceFields
from airflow.models.serialized_dag import SerializedDagModel
@@ -231,6 +232,102 @@ def test_conflicting_task_instance_index_names_are_rejected(dag_maker, bulk):
constructor(**kwargs)
+@pytest.fixture
+def legacy_mapped_ti(dag_maker, session):
+ with dag_maker("legacy_mapping", serialized=True):
+ MockOperator.partial(task_id="mapped").expand(arg2=[1])
+ dr = dag_maker.create_dagrun()
+ session.execute(delete(TI).where(TI.dag_id == dr.dag_id, TI.run_id == dr.run_id))
+ session.execute(delete(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id))
+ ti = TI(
+ dag_maker.serialized_dag.get_task("mapped"),
+ run_id=dr.run_id,
+ dag_version_id=dr.created_dag_version_id,
+ region_index=0,
+ state=TaskInstanceState.SUCCESS,
+ )
+ session.add(ti)
+ session.flush()
+ return dr, ti
+
+
+@pytest.mark.parametrize("whole_task", [False, True])
+def test_legacy_width_one_clear_distinguishes_whole_task(dag_maker, session, legacy_mapped_ti, whole_task):
+ dr, ti = legacy_mapped_ti
+ old_id = ti.id
+
+ dag_maker.serialized_dag.clear(
+ task_ids=["mapped"] if whole_task else [("mapped", 0)], run_id=dr.run_id, session=session
+ )
+ session.flush()
+
+ old = session.get(TI, old_id)
+ assert old.working_set is None
+ assert old.region_id.int == 0
+ assert old.region_index == 0
+ regions = session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)).all()
+ live_query = select(TI).where(TI.dag_id == dr.dag_id, TI.working_set.is_(True))
+ if whole_task:
+ assert old.archived_reason == "superseded"
+ assert len(regions) == 1
+ assert regions[0].node_id == "mapped"
+ assert regions[0].forked_from_region_id is None
+ assert not session.scalars(live_query).all()
+ dr.task_instance_scheduling_decisions(session=session)
+ dr.task_instance_scheduling_decisions(session=session)
+ live = session.scalars(live_query).one()
+ assert live.region_id == regions[0].id
+ assert live.region_index == 0
+ else:
+ assert old.archived_reason == "retry"
+ assert not regions
+ live = session.scalars(live_query).one()
+ assert live.id != old_id
+ assert live.region_id.int == 0
+ assert live.region_index == 0
+
+
+@pytest.mark.parametrize("state", [None, State.SCHEDULED, State.QUEUED, State.UP_FOR_RETRY])
+def test_legacy_whole_clear_archives_pending_execution_in_place(dag_maker, session, legacy_mapped_ti, state):
+ dr, ti = legacy_mapped_ti
+ ti.state = state
+ session.flush()
+ old_id = ti.id
+
+ dag_maker.serialized_dag.clear(task_ids=["mapped"], run_id=dr.run_id, session=session)
+ session.flush()
+
+ archived = session.get(TI, old_id)
+ assert archived.archived_reason == "superseded"
+ assert archived.working_set is None
+ assert archived.state == state
+ assert (archived.start_date, archived.end_date) == (None, None)
+
+
+@pytest.mark.parametrize("replacement_task_id", ["mapped", "different"])
+def test_legacy_reconciliation_does_not_revive_removed_mapping(
+ dag_maker, session, legacy_mapped_ti, replacement_task_id
+):
+ dr, _ = legacy_mapped_ti
+ dag_maker.serialized_dag.clear(task_ids=["mapped"], run_id=dr.run_id, session=session)
+ session.commit()
+ region = session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)).one()
+
+ with dag_maker(dr.dag_id, serialized=True):
+ EmptyOperator(task_id=replacement_task_id)
+ version = DagVersion.get_latest_version(dr.dag_id, session=session)
+ dr.created_dag_version_id = version.id
+ dr.dag = dag_maker.serialized_dag
+ session.flush()
+
+ dr.verify_integrity(dag_version_id=version.id, session=session)
+ dr.verify_integrity(dag_version_id=version.id, session=session)
+ dr.task_instance_scheduling_decisions(session=session)
+
+ assert not session.scalars(select(TI).where(TI.region_id == region.id)).all()
+ assert session.scalars(select(DynamicRegion).where(DynamicRegion.dag_id == dr.dag_id)).all() == [region]
+
+
class TestTaskInstance:
def setup_method(self):
# We don't want to store any code for (test) dags created in this file
diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py
index 3d2734fa0e6..b3284f57adc 100644
--- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py
+++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py
@@ -263,6 +263,17 @@ class TaskIds(RootModel[tuple[str, int]]):
root: Annotated[tuple[str, int], Field(max_length=2, min_length=2)]
+class TaskInstanceIds(RootModel[list[UUID]]):
+ root: Annotated[
+ list[UUID],
+ Field(
+ description="Exact current executions in one DAG run. Mutually exclusive with task_ids and task_group_id.",
+ min_length=1,
+ title="Task Instance Ids",
+ ),
+ ]
+
+
class ClearTaskInstancesBody(BaseModel):
"""
Request body for Clear Task Instances endpoint.
@@ -291,6 +302,27 @@ class ClearTaskInstancesBody(BaseModel):
title="Task Group Id",
),
] = None
+ task_instance_ids: Annotated[
+ TaskInstanceIds | None,
+ Field(
+ description="Exact current executions in one DAG run. Mutually exclusive with task_ids and task_group_id.",
+ title="Task Instance Ids",
+ ),
+ ] = None
+ whole_expansion_ids: Annotated[
+ list[UUID] | None,
+ Field(
+ description="Selected execution IDs whose entire mapped expansion should be cleared.",
+ title="Whole Expansion Ids",
+ ),
+ ] = None
+ include_later_loop_iterations: Annotated[
+ bool | None,
+ Field(
+ description="Clear later iterations when the selection includes a loop gate.",
+ title="Include Later Loop Iterations",
+ ),
+ ] = True
dag_run_id: Annotated[str | None, Field(title="Dag Run Id")] = None
include_upstream: Annotated[bool | None, Field(title="Include Upstream")] = False
include_downstream: Annotated[bool | None, Field(title="Include Downstream")] = False
@@ -1969,6 +2001,8 @@ class BulkTaskInstanceBody(BaseModel):
extra="forbid",
)
new_state: TaskInstanceState | None = None
+ region_id: Annotated[UUID | None, Field(title="Region Id")] = None
+ region_index: Annotated[int | None, Field(title="Region Index")] = None
note: Annotated[Note | None, Field(title="Note")] = None
include_upstream: Annotated[bool | None, Field(title="Include Upstream")] = False
include_downstream: Annotated[bool | None, Field(title="Include Downstream")] = False
@@ -2478,6 +2512,8 @@ class PatchTaskInstanceBody(BaseModel):
extra="forbid",
)
new_state: TaskInstanceState | None = None
+ region_id: Annotated[UUID | None, Field(title="Region Id")] = None
+ region_index: Annotated[int | None, Field(title="Region Index")] = None
note: Annotated[Note | None, Field(title="Note")] = None
include_upstream: Annotated[bool | None, Field(title="Include Upstream")] = False
include_downstream: Annotated[bool | None, Field(title="Include Downstream")] = False
@@ -2622,6 +2658,8 @@ class TaskInstanceResponse(BaseModel):
dag_id: Annotated[str, Field(title="Dag Id")]
dag_run_id: Annotated[str, Field(title="Dag Run Id")]
map_index: Annotated[int, Field(title="Map Index")]
+ region_id: Annotated[UUID, Field(title="Region Id")]
+ region_index: Annotated[int, Field(title="Region Index")]
logical_date: Annotated[datetime | None, Field(title="Logical Date")]
run_after: Annotated[datetime, Field(title="Run After")]
start_date: Annotated[datetime | None, Field(title="Start Date")]
diff --git a/airflow-ctl/tests/airflow_ctl/api/test_operations.py b/airflow-ctl/tests/airflow_ctl/api/test_operations.py
index 122188fc5d2..da959f97460 100644
--- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py
+++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py
@@ -1877,6 +1877,8 @@ class TestTaskInstancesOperations:
dag_id="dag_id",
dag_run_id="dag_run_id",
map_index=-1,
+ region_id=uuid.UUID(int=0),
+ region_index=-1,
run_after=datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc),
state=TaskInstanceState.SUCCESS,
try_number=1,
@@ -1982,6 +1984,8 @@ class TestTasksOperations:
dag_id=dag_id,
dag_run_id="dag_run_id",
map_index=-1,
+ region_id=uuid.UUID(int=0),
+ region_index=-1,
run_after=datetime.datetime(2025, 1, 1, 0, 0, 0),
try_number=1,
max_tries=2,
diff --git a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py
index 5839f87a65f..beb3a024f8c 100644
--- a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py
+++ b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_task_command.py
@@ -61,6 +61,8 @@ class TestFailedDeps:
dag_id=self.dag_id,
dag_run_id=self.run_id,
map_index=-1,
+ region_id=uuid.UUID(int=0),
+ region_index=-1,
logical_date=self.logical_date,
run_after=self.logical_date,
state=state,
@@ -421,6 +423,8 @@ class TestStatesForDagRun:
dag_id=self.dag_id,
dag_run_id=self.run_id,
map_index=map_index,
+ region_id=uuid.UUID(int=0),
+ region_index=map_index,
logical_date=self.logical_date,
run_after=self.logical_date,
state=state,
@@ -672,6 +676,8 @@ class TestState:
dag_id=self.dag_id,
dag_run_id=self.run_id,
map_index=-1,
+ region_id=uuid.UUID(int=0),
+ region_index=-1,
logical_date=self.logical_date,
run_after=self.logical_date,
state=state,
diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
index 7ee007f3fa2..54e5529e7b6 100644
--- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
+++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py
@@ -3882,7 +3882,7 @@ class TestExtractOlInfoFromAssetEvent:
@pytest.mark.db_test
@pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region attribution requires Airflow 3.4")
- def test_cleared_regional_asset_source_keeps_retiring_uuid(self, dag_maker, session):
+ def test_cleared_regional_asset_source_keeps_archiving_uuid(self, dag_maker, session):
with dag_maker(serialized=True):
EmptyOperator(task_id="work")
run = dag_maker.create_dagrun()
@@ -3895,10 +3895,10 @@ class TestExtractOlInfoFromAssetEvent:
ti.region_index = 2
ti.try_number = 1
ti.state = TaskInstanceState.SUCCESS
- retiring_id = ti.id
+ archiving_id = ti.id
event = AssetEvent(
asset_id=asset.id,
- source_task_instance_id=retiring_id,
+ source_task_instance_id=archiving_id,
source_dag_id=ti.dag_id,
source_task_id=ti.task_id,
source_run_id=ti.run_id,
@@ -3907,7 +3907,7 @@ class TestExtractOlInfoFromAssetEvent:
session.add(event)
session.flush()
successor = ti.prepare_db_for_next_try(session)
- assert successor.id != retiring_id
+ assert successor.id != archiving_id
session.commit()
session.expire(event)
assert event.uri == asset.uri
@@ -3916,7 +3916,7 @@ class TestExtractOlInfoFromAssetEvent:
dependencies = _get_ol_job_dependencies_from_asset_events([event] * 10)
assert len(dependencies) == 1
- assert dependencies[0].get("run_id") == str(retiring_id)
+ assert dependencies[0].get("run_id") == str(archiving_id)
assert dependencies[0]["job_name"] == f"{run.dag_id}.work"
assert len(dependencies[0]["asset_events"]) == 10
assert ti in session