Skip to content

Commit d4aa44e

Browse files
committed
Add durable task instance launch records
Persist immutable executor launch tokens independently of mutable or deleted TaskInstance rows so that expected stale executor launches return a typed 409 stale_executor_launch instead of an opaque 404. Refs #69760
1 parent 3968cf0 commit d4aa44e

32 files changed

Lines changed: 1098 additions & 21 deletions

File tree

airflow-core/docs/migrations-ref.rst

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are executed via when you ru
3939
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
4040
| Revision ID | Revises ID | Airflow Version | Description |
4141
+=========================+==================+===================+==============================================================+
42-
| ``c7f0a5d2e9b4`` (head) | ``76c46545c91e`` | ``3.4.0`` | Lower case team names. |
42+
| ``3c5f8e9a1d2b`` (head) | ``c7f0a5d2e9b4`` | ``3.4.0`` | Add task_instance_launch table for durable executor launch |
43+
| | | | records. |
44+
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
45+
| ``c7f0a5d2e9b4`` | ``76c46545c91e`` | ``3.4.0`` | Lower case team names. |
4346
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
4447
| ``76c46545c91e`` | ``3c525f44bea8`` | ``3.4.0`` | Add new index for trigger. |
4548
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
Durable executor launch records fence stale task executors
2+
3+
Executor task launches are now recorded in a new ``task_instance_launch`` table that is
4+
independent of the mutable ``task_instance`` row. The scheduler pre-assigns an immutable
5+
launch token when a task instance moves from ``scheduled`` to ``queued`` and writes an
6+
``active`` launch record for it. When a worker calls the Execution API ``/run`` endpoint,
7+
the token it carries is validated against the durable record: a launch whose token has
8+
already reached a terminal state (``consumed`` or ``superseded``) now returns a typed
9+
``409 stale_executor_launch`` response instead of a generic ``404``, so a worker started
10+
from a stale/duplicated launch exits cleanly (return code 0) rather than corrupting the
11+
task instance's state. Launch records are superseded on reschedule of stuck-queued tasks,
12+
failed pod adoption, orphan reset, and task-instance clear/next-try; successful adoption
13+
preserves the token. The mechanism is executor-agnostic (CeleryExecutor, KubernetesExecutor,
14+
EdgeExecutor).
15+
16+
The new ``409 stale_executor_launch`` behaviour is gated behind a new Execution API version
17+
(``v2026_06_30``); older Task SDK clients pinned to earlier Execution API versions continue
18+
to observe the previous ``404`` behaviour, so the change is backward compatible.
19+
20+
* Types of change
21+
22+
* [ ] Dag changes
23+
* [ ] Config changes
24+
* [x] API changes
25+
* [ ] CLI changes
26+
* [x] Behaviour changes
27+
* [ ] Plugin changes
28+
* [ ] Dependency changes
29+
* [ ] Code interface changes

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,8 @@ class TIEnterRunningPayload(StrictBaseModel):
6565
"""Process Identifier on `hostname`"""
6666
start_date: UtcDateTime
6767
"""When the task started executing"""
68+
external_executor_id: str | None = None
69+
"""Executor token for durable launch validation"""
6870

6971

7072
# Create an enum to give a nice name in the generated datamodels
@@ -287,6 +289,8 @@ class TaskInstance(BaseModel):
287289
map_index: int = -1
288290
hostname: str | None = None
289291
context_carrier: dict | None = None
292+
external_executor_id: str | None = None
293+
"""Executor token for durable launch validation"""
290294
# The supervisor routes tasks to a coordinator by queue. The default keeps
291295
# hand-built instances (tests, dry runs) valid; the executor workload
292296
# always sends the real value.

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@
8989
from airflow.models.dagrun import DagRun as DR
9090
from airflow.models.hitl import HITLDetail
9191
from airflow.models.log import Log
92+
from airflow.models.task_instance_launch import TERMINAL_STATES, TaskInstanceLaunch
9293
from airflow.models.taskinstance import TaskInstance as TI, _stop_remaining_tasks
9394
from airflow.models.taskinstancehistory import TaskInstanceHistory as TIH
9495
from airflow.models.taskreschedule import TaskReschedule
@@ -170,6 +171,7 @@ def ti_run(
170171
TI.unixname,
171172
TI.pid,
172173
TI.dag_version_id,
174+
TI.external_executor_id,
173175
# This selects the raw JSON value, bypassing the deserialization -- we want that to happen on the
174176
# client
175177
column("next_kwargs", JSON),
@@ -186,6 +188,25 @@ def ti_run(
186188
ti = session.execute(old).one()
187189
log.debug("Retrieved task instance details", state=ti.state, dag_id=ti.dag_id, task_id=ti.task_id)
188190
except NoResultFound:
191+
# TI not found - check if we have a token to validate against launch records
192+
payload_token = ti_run_payload.external_executor_id
193+
if payload_token:
194+
# Check if this token has a known launch record
195+
launch = TaskInstanceLaunch.get_by_token(payload_token, session=session)
196+
if launch and launch.state in TERMINAL_STATES:
197+
# Token is known but in a terminal state (consumed or superseded)
198+
log.warning(
199+
"Stale executor token for missing task instance", token=payload_token, state=launch.state
200+
)
201+
raise HTTPException(
202+
status_code=status.HTTP_409_CONFLICT,
203+
detail={
204+
"reason": "stale_executor_launch",
205+
"message": "Executor token is stale or has been superseded",
206+
},
207+
)
208+
# Token is unknown/null - ambiguous, could be upgrade scenario or stale executor
209+
# Let the 404 be returned
189210
log.error("Task Instance not found")
190211
raise HTTPException(
191212
status_code=status.HTTP_404_NOT_FOUND,
@@ -195,6 +216,39 @@ def ti_run(
195216
},
196217
)
197218

219+
# Validate launch token if provided (launch record presence drives enforcement)
220+
# If TI has a DB token and that token has a launch record, enforce strict matching
221+
payload_token = ti_run_payload.external_executor_id
222+
db_token = ti.external_executor_id
223+
224+
if db_token:
225+
# Check if this DB token has an active/terminal launch record
226+
launch_record = TaskInstanceLaunch.get_by_token(db_token, session=session)
227+
if launch_record:
228+
# Launch record exists - enforce strict matching
229+
if payload_token != db_token:
230+
log.error(
231+
"External executor ID mismatch for known launch record",
232+
payload_token=payload_token,
233+
db_token=db_token,
234+
launch_state=launch_record.state,
235+
)
236+
raise HTTPException(
237+
status_code=status.HTTP_409_CONFLICT,
238+
detail={
239+
"reason": "stale_executor_launch",
240+
"message": "Executor token does not match known launch record",
241+
},
242+
)
243+
elif payload_token and payload_token != db_token:
244+
# No launch record exists, preserve rolling compatibility
245+
# Only enforce if both have tokens and they differ
246+
log.debug(
247+
"Token mismatch but no launch record (rolling upgrade scenario)",
248+
payload_token=payload_token,
249+
db_token=db_token,
250+
)
251+
198252
# We exclude_unset to avoid updating fields that are not set in the payload
199253
data = ti_run_payload.model_dump(exclude_unset=True)
200254

@@ -265,6 +319,32 @@ def ti_run(
265319
result = session.execute(query)
266320
log.info("Task instance state updated", rows_affected=getattr(result, "rowcount", 0))
267321

322+
# Mark the launch record as consumed if we have a token and the transition to RUNNING succeeded.
323+
# Only mark on successful transitions from QUEUED/RESTARTING (not on duplicate RUNNING requests).
324+
if payload_token and previous_state in (TaskInstanceState.QUEUED, TaskInstanceState.RESTARTING):
325+
marked = TaskInstanceLaunch.mark_consumed(
326+
token=payload_token,
327+
session=session,
328+
)
329+
if marked:
330+
log.debug("Marked launch token as consumed", token=payload_token)
331+
else:
332+
# If mark_consumed returns False, the record was already consumed or missing
333+
# Only allow on duplicate RUNNING request (same host/user/pid), otherwise it's a conflict
334+
if previous_state == TaskInstanceState.RUNNING:
335+
log.debug("Launch token already consumed (duplicate request)", token=payload_token)
336+
else:
337+
# Roll back the update since mark_consumed failed and this isn't a duplicate
338+
session.rollback()
339+
log.error("Failed to mark launch token as consumed", token=payload_token)
340+
raise HTTPException(
341+
status_code=status.HTTP_409_CONFLICT,
342+
detail={
343+
"reason": "stale_executor_launch",
344+
"message": "Launch record not found or already consumed",
345+
},
346+
)
347+
268348
dr = (
269349
session.scalars(
270350
select(DR)

airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
AddPartitionDateField,
4848
AddRetryPolicyFields,
4949
AddTaskAndAssetStateStoreEndpoints,
50+
AddTaskInstanceExternalExecutorIdField,
5051
AddTaskInstanceQueueField,
5152
AddTeamNameField,
5253
AddVariableKeysEndpoint,
@@ -61,6 +62,7 @@
6162
AddVariableKeysEndpoint,
6263
AddConnectionTestEndpoint,
6364
AddAwaitingInputStatePayload,
65+
AddTaskInstanceExternalExecutorIdField,
6466
AddTaskInstanceQueueField,
6567
AddRetryPolicyFields,
6668
AddTeamNameField,

airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
DagRun,
3030
TaskInstance,
3131
TIAwaitingInputStatePayload,
32+
TIEnterRunningPayload,
3233
TIRetryStatePayload,
3334
TIRunContext,
3435
)
@@ -61,6 +62,17 @@ class AddTaskInstanceQueueField(VersionChange):
6162
instructions_to_migrate_to_previous_version = (schema(TaskInstance).field("queue").didnt_exist,)
6263

6364

65+
class AddTaskInstanceExternalExecutorIdField(VersionChange):
66+
"""Add the `external_executor_id` field to task instance launch payloads."""
67+
68+
description = __doc__
69+
70+
instructions_to_migrate_to_previous_version = (
71+
schema(TaskInstance).field("external_executor_id").didnt_exist,
72+
schema(TIEnterRunningPayload).field("external_executor_id").didnt_exist,
73+
)
74+
75+
6476
class AddAwaitingInputStatePayload(VersionChange):
6577
"""Add the awaiting_input task instance state transition payload (Human-in-the-loop, no trigger)."""
6678

airflow-core/src/airflow/executors/workloads/task.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,12 @@ class TaskInstanceDTO(TaskInstance):
4545
pool_slots: int
4646
priority_weight: int
4747

48-
external_executor_id: str | None = Field(default=None, exclude=True)
48+
# NOT excluded: the durable launch token must reach the worker so it can echo
49+
# it back on POST /run for stale-executor fencing. Executors that pre-assign
50+
# external_executor_id (Celery, Kubernetes) rely on this surviving serialization
51+
# into the workload JSON that is shipped to the worker/pod. executor_config stays
52+
# excluded — it is a scheduler-only concern the worker never reads.
53+
external_executor_id: str | None = Field(default=None)
4954
executor_config: dict | None = Field(default=None, exclude=True)
5055

5156
# TODO: Task-SDK: Can we replace TaskInstanceKey with just the uuid across the codebase?

airflow-core/src/airflow/jobs/scheduler_job_runner.py

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@
105105
from airflow.models.dagwarning import DagWarning, DagWarningType
106106
from airflow.models.pool import normalize_pool_name_for_stats
107107
from airflow.models.serialized_dag import SerializedDagModel
108+
from airflow.models.task_instance_launch import TaskInstanceLaunch, TaskInstanceLaunchState
108109
from airflow.models.taskinstance import TaskInstance
109110
from airflow.models.taskinstancekey import TaskInstanceKey
110111
from airflow.models.team import Team
@@ -1117,6 +1118,30 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
11171118
else:
11181119
session.execute(queued_update)
11191120

1121+
# Bulk insert TaskInstanceLaunch records for active launches with non-null tokens.
1122+
# This happens after external_executor_ids have been synced into executable_tis in-memory.
1123+
# The DB-stored external_executor_id values will be read back in the next query if needed.
1124+
launch_records = []
1125+
for ti in executable_tis:
1126+
token = ti.external_executor_id
1127+
if token:
1128+
launch_records.append(
1129+
TaskInstanceLaunch(
1130+
token=token,
1131+
task_instance_id=str(ti.id),
1132+
dag_id=ti.dag_id,
1133+
task_id=ti.task_id,
1134+
run_id=ti.run_id,
1135+
map_index=ti.map_index if ti.map_index is not None else -1,
1136+
try_number=ti.try_number,
1137+
executor=ti.executor or "default",
1138+
state=TaskInstanceLaunchState.ACTIVE,
1139+
)
1140+
)
1141+
if launch_records:
1142+
session.add_all(launch_records)
1143+
session.flush()
1144+
11201145
for ti in executable_tis:
11211146
ti.emit_state_change_metric(TaskInstanceState.QUEUED)
11221147

@@ -3212,6 +3237,15 @@ def _reschedule_stuck_task(self, ti: TaskInstance, session: Session):
32123237
filter_for_tis = TI.filter_for_tis([ti])
32133238
if filter_for_tis is None:
32143239
return
3240+
3241+
# Supersede any active launch record before rescheduling, so the old token
3242+
# cannot be reused if the old executor manages to come back.
3243+
if ti.external_executor_id:
3244+
TaskInstanceLaunch.mark_superseded(
3245+
token=ti.external_executor_id,
3246+
session=session,
3247+
)
3248+
32153249
session.execute(
32163250
update(TI)
32173251
.where(filter_for_tis)
@@ -3220,6 +3254,7 @@ def _reschedule_stuck_task(self, ti: TaskInstance, session: Session):
32203254
queued_dttm=None,
32213255
queued_by_job_id=None,
32223256
scheduled_dttm=timezone.utcnow(),
3257+
external_executor_id=None, # Clear so next queue cycle assigns fresh token
32233258
)
32243259
.execution_options(synchronize_session=False)
32253260
)
@@ -3428,6 +3463,13 @@ def adopt_or_reset_orphaned_tasks(self, *, session: Session = NEW_SESSION) -> in
34283463
reset_tis_message = []
34293464
for ti in to_reset:
34303465
reset_tis_message.append(repr(ti))
3466+
# Mark the current launch token as superseded before resetting the TI
3467+
# so the old executor cannot re-adopt this task if it comes back
3468+
if ti.external_executor_id:
3469+
TaskInstanceLaunch.mark_superseded(
3470+
token=ti.external_executor_id,
3471+
session=session,
3472+
)
34313473
# If we reset a TI, it will be eligible to be scheduled again.
34323474
# This can cause the scheduler to increase the try_number on the TI.
34333475
# Record the current try to TaskInstanceHistory first so users have an audit trail for
Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
19+
"""
20+
Add task_instance_launch table for durable executor launch records.
21+
22+
Revision ID: 3c5f8e9a1d2b
23+
Revises: c7f0a5d2e9b4
24+
Create Date: 2026-05-10 00:00:00.000000
25+
26+
"""
27+
28+
from __future__ import annotations
29+
30+
import sqlalchemy as sa
31+
from alembic import op
32+
33+
from airflow.utils.sqlalchemy import UtcDateTime
34+
35+
# revision identifiers, used by Alembic.
36+
revision = "3c5f8e9a1d2b"
37+
down_revision = "c7f0a5d2e9b4"
38+
branch_labels = None
39+
depends_on = None
40+
airflow_version = "3.4.0"
41+
42+
43+
def upgrade():
44+
"""Create task_instance_launch table for durable executor token tracking."""
45+
op.create_table(
46+
"task_instance_launch",
47+
sa.Column("token", sa.String(256), nullable=False),
48+
sa.Column("task_instance_id", sa.String(250), nullable=False),
49+
sa.Column("dag_id", sa.String(250), nullable=False),
50+
sa.Column("task_id", sa.String(250), nullable=False),
51+
sa.Column("run_id", sa.String(250), nullable=False),
52+
sa.Column("map_index", sa.Integer(), nullable=False, server_default="-1"),
53+
sa.Column("try_number", sa.Integer(), nullable=False),
54+
sa.Column("executor", sa.String(256), nullable=False),
55+
sa.Column("state", sa.String(20), nullable=False, server_default="active"),
56+
sa.Column("created_at", UtcDateTime(timezone=True), nullable=False, server_default=sa.func.now()),
57+
sa.Column("updated_at", UtcDateTime(timezone=True), nullable=False, server_default=sa.func.now()),
58+
sa.Column("consumed_at", UtcDateTime(timezone=True), nullable=True),
59+
sa.Column("superseded_at", UtcDateTime(timezone=True), nullable=True),
60+
sa.PrimaryKeyConstraint("token", name="pk_task_instance_launch_token"),
61+
sa.CheckConstraint(
62+
"state IN ('active', 'consumed', 'superseded')",
63+
name="state_enum",
64+
),
65+
)
66+
op.create_index(
67+
"idx_task_instance_launch_task_instance_id",
68+
"task_instance_launch",
69+
["task_instance_id"],
70+
)
71+
op.create_index(
72+
"idx_task_instance_launch_state_updated",
73+
"task_instance_launch",
74+
["state", "updated_at"],
75+
)
76+
op.create_index(
77+
"idx_task_instance_launch_created_at",
78+
"task_instance_launch",
79+
["created_at"],
80+
)
81+
82+
83+
def downgrade():
84+
"""Drop task_instance_launch table."""
85+
op.drop_table("task_instance_launch")

0 commit comments

Comments
 (0)