Skip to content

Commit 56f29f5

Browse files
ashbashb
committed
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.
1 parent 0d3ae91 commit 56f29f5

36 files changed

Lines changed: 5371 additions & 266 deletions

File tree

‎airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
)
3636

3737
from airflow._shared.secrets_masker import redact
38+
from airflow.api_fastapi.common.region import OmitsMissingRegion, RegionId, RegionIndex
3839
from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel
3940
from airflow.api_fastapi.core_api.datamodels.dag_versions import DagVersionResponse
4041
from airflow.api_fastapi.core_api.datamodels.job import JobResponse
@@ -49,14 +50,16 @@ class NewTaskResponse(BaseModel):
4950
task_display_name: str
5051

5152

52-
class TaskInstanceResponse(BaseModel):
53+
class TaskInstanceResponse(OmitsMissingRegion, BaseModel):
5354
"""TaskInstance serializer for responses."""
5455

5556
id: UUID
5657
task_id: str
5758
dag_id: str
5859
run_id: str = Field(alias="dag_run_id")
5960
map_index: int
61+
region_id: RegionId = None
62+
region_index: RegionIndex = None
6063
logical_date: datetime | None
6164
run_after: datetime
6265
start_date: datetime | None
@@ -239,6 +242,19 @@ class ClearTaskInstancesBody(StrictBaseModel):
239242
"The group's tasks are resolved on the server from the dag structure, so all of them are "
240243
"targeted regardless of how many there are.",
241244
)
245+
task_instance_ids: list[UUID] | None = Field(
246+
default=None,
247+
min_length=1,
248+
description="Exact current executions in one DAG run. Mutually exclusive with task_ids and task_group_id.",
249+
)
250+
whole_expansion_ids: list[UUID] = Field(
251+
default_factory=list,
252+
description="Selected execution IDs whose entire mapped expansion should be cleared.",
253+
)
254+
include_later_loop_iterations: bool = Field(
255+
default=True,
256+
description="Clear later iterations when the selection includes a loop gate.",
257+
)
242258
dag_run_id: str | None = None
243259
include_upstream: bool = False
244260
include_downstream: bool = False
@@ -281,11 +297,26 @@ def validate_model(cls, data: Any) -> Any:
281297
raise ValueError("Only one of task_ids or task_group_id may be provided")
282298
return data
283299

300+
@model_validator(mode="after")
301+
def validate_execution_selection(self):
302+
if self.task_instance_ids is not None:
303+
if self.task_ids is not None or self.task_group_id is not None:
304+
raise ValueError("task_instance_ids cannot be combined with task_ids or task_group_id")
305+
if self.dag_run_id is None or self.include_past or self.include_future:
306+
raise ValueError(
307+
"task_instance_ids requires one dag_run_id without include_past or include_future"
308+
)
309+
if not set(self.whole_expansion_ids) <= set(self.task_instance_ids or ()):
310+
raise ValueError("whole_expansion_ids must be a subset of task_instance_ids")
311+
return self
312+
284313

285314
class PatchTaskInstanceBody(StrictBaseModel):
286315
"""Request body for patching task instance state."""
287316

288317
new_state: TaskInstanceState | None = None
318+
region_id: UUID | None = None
319+
region_index: int | None = None
289320
note: Annotated[str, StringConstraints(max_length=1000)] | None = None
290321
include_upstream: bool = False
291322
include_downstream: bool = False

‎airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4975,6 +4975,13 @@ components:
49754975
map_index:
49764976
type: integer
49774977
title: Map Index
4978+
region_id:
4979+
type: string
4980+
format: uuid
4981+
title: Region Id
4982+
region_index:
4983+
type: integer
4984+
title: Region Index
49784985
logical_date:
49794986
anyOf:
49804987
- type: string

‎airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7540,6 +7540,23 @@ paths:
75407540
type: string
75417541
- type: 'null'
75427542
title: Update Mask
7543+
- name: region_id
7544+
in: query
7545+
required: false
7546+
schema:
7547+
anyOf:
7548+
- type: string
7549+
format: uuid
7550+
- type: 'null'
7551+
title: Region Id
7552+
- name: region_index
7553+
in: query
7554+
required: false
7555+
schema:
7556+
anyOf:
7557+
- type: integer
7558+
- type: 'null'
7559+
title: Region Index
75437560
requestBody:
75447561
required: true
75457562
content:
@@ -8505,6 +8522,23 @@ paths:
85058522
type: string
85068523
- type: 'null'
85078524
title: Update Mask
8525+
- name: region_id
8526+
in: query
8527+
required: false
8528+
schema:
8529+
anyOf:
8530+
- type: string
8531+
format: uuid
8532+
- type: 'null'
8533+
title: Region Id
8534+
- name: region_index
8535+
in: query
8536+
required: false
8537+
schema:
8538+
anyOf:
8539+
- type: integer
8540+
- type: 'null'
8541+
title: Region Index
85088542
requestBody:
85098543
required: true
85108544
content:
@@ -9543,6 +9577,23 @@ paths:
95439577
type: string
95449578
- type: 'null'
95459579
title: Update Mask
9580+
- name: region_id
9581+
in: query
9582+
required: false
9583+
schema:
9584+
anyOf:
9585+
- type: string
9586+
format: uuid
9587+
- type: 'null'
9588+
title: Region Id
9589+
- name: region_index
9590+
in: query
9591+
required: false
9592+
schema:
9593+
anyOf:
9594+
- type: integer
9595+
- type: 'null'
9596+
title: Region Index
95469597
requestBody:
95479598
required: true
95489599
content:
@@ -9621,6 +9672,23 @@ paths:
96219672
schema:
96229673
type: string
96239674
title: Group Id
9675+
- name: region_id
9676+
in: query
9677+
required: false
9678+
schema:
9679+
anyOf:
9680+
- type: string
9681+
format: uuid
9682+
- type: 'null'
9683+
title: Region Id
9684+
- name: region_index
9685+
in: query
9686+
required: false
9687+
schema:
9688+
anyOf:
9689+
- type: integer
9690+
- type: 'null'
9691+
title: Region Index
96249692
requestBody:
96259693
required: true
96269694
content:
@@ -9658,6 +9726,12 @@ paths:
96589726
schema:
96599727
$ref: '#/components/schemas/HTTPExceptionResponse'
96609728
description: Not Found
9729+
'409':
9730+
content:
9731+
application/json:
9732+
schema:
9733+
$ref: '#/components/schemas/HTTPExceptionResponse'
9734+
description: Conflict
96619735
'422':
96629736
description: Validation Error
96639737
content:
@@ -9711,6 +9785,23 @@ paths:
97119785
type: string
97129786
- type: 'null'
97139787
title: Update Mask
9788+
- name: region_id
9789+
in: query
9790+
required: false
9791+
schema:
9792+
anyOf:
9793+
- type: string
9794+
format: uuid
9795+
- type: 'null'
9796+
title: Region Id
9797+
- name: region_index
9798+
in: query
9799+
required: false
9800+
schema:
9801+
anyOf:
9802+
- type: integer
9803+
- type: 'null'
9804+
title: Region Index
97149805
requestBody:
97159806
required: true
97169807
content:
@@ -9748,6 +9839,12 @@ paths:
97489839
schema:
97499840
$ref: '#/components/schemas/HTTPExceptionResponse'
97509841
description: Not Found
9842+
'409':
9843+
content:
9844+
application/json:
9845+
schema:
9846+
$ref: '#/components/schemas/HTTPExceptionResponse'
9847+
description: Conflict
97519848
'422':
97529849
description: Validation Error
97539850
content:
@@ -9801,6 +9898,23 @@ paths:
98019898
type: string
98029899
- type: 'null'
98039900
title: Update Mask
9901+
- name: region_id
9902+
in: query
9903+
required: false
9904+
schema:
9905+
anyOf:
9906+
- type: string
9907+
format: uuid
9908+
- type: 'null'
9909+
title: Region Id
9910+
- name: region_index
9911+
in: query
9912+
required: false
9913+
schema:
9914+
anyOf:
9915+
- type: integer
9916+
- type: 'null'
9917+
title: Region Index
98049918
requestBody:
98059919
required: true
98069920
content:
@@ -9838,6 +9952,12 @@ paths:
98389952
schema:
98399953
$ref: '#/components/schemas/HTTPExceptionResponse'
98409954
description: Not Found
9955+
'409':
9956+
content:
9957+
application/json:
9958+
schema:
9959+
$ref: '#/components/schemas/HTTPExceptionResponse'
9960+
description: Conflict
98419961
'422':
98429962
description: Validation Error
98439963
content:
@@ -12827,6 +12947,17 @@ components:
1282712947
anyOf:
1282812948
- $ref: '#/components/schemas/TaskInstanceState'
1282912949
- type: 'null'
12950+
region_id:
12951+
anyOf:
12952+
- type: string
12953+
format: uuid
12954+
- type: 'null'
12955+
title: Region Id
12956+
region_index:
12957+
anyOf:
12958+
- type: integer
12959+
- type: 'null'
12960+
title: Region Index
1283012961
note:
1283112962
anyOf:
1283212963
- type: string
@@ -13206,6 +13337,30 @@ components:
1320613337
description: Clear every task in this task group. Mutually exclusive with
1320713338
`task_ids`. The group's tasks are resolved on the server from the dag
1320813339
structure, so all of them are targeted regardless of how many there are.
13340+
task_instance_ids:
13341+
anyOf:
13342+
- items:
13343+
type: string
13344+
format: uuid
13345+
type: array
13346+
minItems: 1
13347+
- type: 'null'
13348+
title: Task Instance Ids
13349+
description: Exact current executions in one DAG run. Mutually exclusive
13350+
with task_ids and task_group_id.
13351+
whole_expansion_ids:
13352+
items:
13353+
type: string
13354+
format: uuid
13355+
type: array
13356+
title: Whole Expansion Ids
13357+
description: Selected execution IDs whose entire mapped expansion should
13358+
be cleared.
13359+
include_later_loop_iterations:
13360+
type: boolean
13361+
title: Include Later Loop Iterations
13362+
description: Clear later iterations when the selection includes a loop gate.
13363+
default: true
1320913364
dag_run_id:
1321013365
anyOf:
1321113366
- type: string
@@ -16115,6 +16270,17 @@ components:
1611516270
anyOf:
1611616271
- $ref: '#/components/schemas/TaskInstanceState'
1611716272
- type: 'null'
16273+
region_id:
16274+
anyOf:
16275+
- type: string
16276+
format: uuid
16277+
- type: 'null'
16278+
title: Region Id
16279+
region_index:
16280+
anyOf:
16281+
- type: integer
16282+
- type: 'null'
16283+
title: Region Index
1611816284
note:
1611916285
anyOf:
1612016286
- type: string
@@ -16964,6 +17130,13 @@ components:
1696417130
map_index:
1696517131
type: integer
1696617132
title: Map Index
17133+
region_id:
17134+
type: string
17135+
format: uuid
17136+
title: Region Id
17137+
region_index:
17138+
type: integer
17139+
title: Region Index
1696717140
logical_date:
1696817141
anyOf:
1696917142
- type: string

0 commit comments

Comments
 (0)