Validate KubernetesExecutor launches with external executor IDs

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

@@ -64,6 +64,8 @@ class TIEnterRunningPayload(StrictBaseModel):
     """Process Identifier on `hostname`"""
     start_date: UtcDateTime
     """When the task started executing"""
+    external_executor_id: str | None = None
+    """Executor launch token assigned when the task was queued"""
 
 
 # Create an enum to give a nice name in the generated datamodels
@@ -290,6 +292,7 @@ class TaskInstance(BaseModel):
     # hand-built instances (tests, dry runs) valid; the executor workload
     # always sends the real value.
     queue: str = "default"
+    external_executor_id: str | None = None
 
 
 class AssetReferenceAssetEventDagRun(StrictBaseModel):

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

@@ -163,6 +163,7 @@ def ti_run(
             TI.hostname,
             TI.unixname,
             TI.pid,
+            TI.external_executor_id,
             # This selects the raw JSON value, bypassing the deserialization -- we want that to happen on the
             # client
             column("next_kwargs", JSON),
@@ -190,6 +191,13 @@ def ti_run(
 
     # We exclude_unset to avoid updating fields that are not set in the payload
     data = ti_run_payload.model_dump(exclude_unset=True)
+    # Only fence on the launch token when the worker actually presented one. Older Task SDK
+    # clients (e.g. mid rolling-upgrade, before this field existed) omit it entirely; treating
+    # that as a stale launch would 409 every task start for pre-assigning executors
+    # (KubernetesExecutor, CeleryExecutor) until every worker is upgraded. When it is absent we
+    # fall back to state-based validation, matching the pre-token behavior.
+    payload_has_external_executor_id = "external_executor_id" in data
+    payload_external_executor_id = data.pop("external_executor_id", None)
 
     # don't update start date when resuming from deferral
     if ti.next_kwargs:
@@ -208,6 +216,23 @@ def ti_run(
         ti_run_payload.pid,
     ):
         log.info("Duplicate start request received", hostname=ti_run_payload.hostname)
+    elif (
+        payload_has_external_executor_id
+        and ti.external_executor_id is not None
+        and ti.external_executor_id != payload_external_executor_id
+    ):
+        log.warning(
+            "Cannot start Task Instance with stale executor launch token",
+            expected_external_executor_id=ti.external_executor_id,
+            provided_external_executor_id=payload_external_executor_id,
+        )
+        raise HTTPException(
+            status_code=status.HTTP_409_CONFLICT,
+            detail={
+                "reason": "stale_executor_launch",
+                "message": "TI executor launch token does not match the current queued task instance",
+            },
+        )
     elif previous_state not in (TaskInstanceState.QUEUED, TaskInstanceState.RESTARTING):
         log.warning(
             "Cannot start Task Instance in invalid state",

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

@@ -45,7 +45,7 @@ class TaskInstanceDTO(TaskInstance):
     pool_slots: int
     priority_weight: int
 
-    external_executor_id: str | None = Field(default=None, exclude=True)
+    external_executor_id: str | None = None
     executor_config: dict | None = Field(default=None, exclude=True)
 
     # TODO: Task-SDK: Can we replace TaskInstanceKey with just the uuid across the codebase?

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py

@@ -322,6 +322,95 @@ def test_ti_run_state_to_running(
         )
         assert response.status_code == 409
 
+    def test_ti_run_rejects_stale_external_executor_id(self, client, session, create_task_instance):
+        """A worker can only start the launch token currently assigned to the queued TI."""
+        ti = create_task_instance(
+            task_id="test_ti_run_rejects_stale_external_executor_id",
+            state=State.QUEUED,
+            dagrun_state=DagRunState.RUNNING,
+            session=session,
+            dag_id=str(uuid4()),
+        )
+        ti.external_executor_id = "current-launch-token"
+        session.commit()
+
+        response = client.patch(
+            f"/execution/task-instances/{ti.id}/run",
+            json={
+                "state": "running",
+                "hostname": "random-hostname",
+                "unixname": "random-unixname",
+                "pid": 100,
+                "start_date": "2024-09-30T12:00:00Z",
+                "external_executor_id": "stale-launch-token",
+            },
+        )
+
+        assert response.status_code == 409
+        assert response.json()["detail"]["reason"] == "stale_executor_launch"
+        session.refresh(ti)
+        assert ti.state == State.QUEUED
+
+    def test_ti_run_accepts_matching_external_executor_id(self, client, session, create_task_instance):
+        """A worker presenting the launch token currently assigned to the TI can start it."""
+        ti = create_task_instance(
+            task_id="test_ti_run_accepts_matching_external_executor_id",
+            state=State.QUEUED,
+            dagrun_state=DagRunState.RUNNING,
+            session=session,
+            dag_id=str(uuid4()),
+        )
+        ti.external_executor_id = "current-launch-token"
+        session.commit()
+
+        response = client.patch(
+            f"/execution/task-instances/{ti.id}/run",
+            json={
+                "state": "running",
+                "hostname": "random-hostname",
+                "unixname": "random-unixname",
+                "pid": 100,
+                "start_date": "2024-09-30T12:00:00Z",
+                "external_executor_id": "current-launch-token",
+            },
+        )
+
+        assert response.status_code == 200
+        session.refresh(ti)
+        assert ti.state == State.RUNNING
+        # The launch token is preserved; matching it must not clear or overwrite the field.
+        assert ti.external_executor_id == "current-launch-token"
+
+    def test_ti_run_allows_missing_external_executor_id_for_old_clients(
+        self, client, session, create_task_instance
+    ):
+        """An older worker that omits the launch token must not be fenced out (rolling upgrade)."""
+        ti = create_task_instance(
+            task_id="test_ti_run_allows_missing_external_executor_id_for_old_clients",
+            state=State.QUEUED,
+            dagrun_state=DagRunState.RUNNING,
+            session=session,
+            dag_id=str(uuid4()),
+        )
+        ti.external_executor_id = "current-launch-token"
+        session.commit()
+
+        # Payload deliberately omits ``external_executor_id`` entirely, as a pre-token Task SDK would.
+        response = client.patch(
+            f"/execution/task-instances/{ti.id}/run",
+            json={
+                "state": "running",
+                "hostname": "random-hostname",
+                "unixname": "random-unixname",
+                "pid": 100,
+                "start_date": "2024-09-30T12:00:00Z",
+            },
+        )
+
+        assert response.status_code == 200
+        session.refresh(ti)
+        assert ti.state == State.RUNNING
+
     def test_ti_run_returns_execution_token(
         self, client, exec_app, session, create_task_instance, time_machine
     ):

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py

@@ -33,7 +33,7 @@
 from contextlib import suppress
 from datetime import datetime, timedelta
 from queue import Empty, Queue
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from deprecated import deprecated
 from kubernetes.dynamic import DynamicClient
@@ -81,6 +81,7 @@ class KubernetesExecutor(BaseExecutor):
     RUNNING_POD_LOG_LINES = 100
     supports_ad_hoc_ti_run: bool = True
     supports_multi_team: bool = True
+    pre_assigns_external_executor_id: ClassVar[bool] = True
 
     if TYPE_CHECKING and AIRFLOW_V_3_0_PLUS:
         # In the v3 path, we store workloads, not commands as strings.
@@ -318,7 +319,16 @@ def execute_async(
                     queue,
                 )
 
-        self.event_buffer[key] = (TaskInstanceState.QUEUED, self.scheduler_job_id)
+        event_info = self.scheduler_job_id
+        try:
+            from airflow.executors.workloads import ExecuteTask
+
+            if len(command) == 1 and isinstance(command[0], ExecuteTask):
+                event_info = command[0].ti.external_executor_id or self.scheduler_job_id
+        except TypeError:
+            pass
+
+        self.event_buffer[key] = (TaskInstanceState.QUEUED, event_info)
         self.task_queue.put(
             KubernetesJob(key, command, kube_executor_config, pod_template_file, coordinator_kube_image)
         )
@@ -349,6 +359,91 @@ def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
             self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config)
             self.running.add(key)
 
+    @provide_session
+    def _should_create_pod_for_job(self, task: KubernetesJob, *, session: Session = NEW_SESSION) -> bool:
+        """Check whether a queued Kubernetes job still owns the current task launch token."""
+        from airflow.executors.workloads import ExecuteTask
+        from airflow.models.taskinstance import TaskInstance
+
+        if not task.command or not isinstance(task.command[0], ExecuteTask):
+            return True
+
+        workload = task.command[0]
+        workload_ti = workload.ti
+        launch_token = workload_ti.external_executor_id
+
+        try:
+            scheduler_job_id = int(self.scheduler_job_id) if self.scheduler_job_id is not None else None
+        except ValueError:
+            self.log.debug(
+                "Skipping stale Kubernetes workload check because scheduler_job_id %r is not numeric",
+                self.scheduler_job_id,
+            )
+            return True
+
+        ti = session.execute(
+            select(
+                TaskInstance.id,
+                TaskInstance.state,
+                TaskInstance.try_number,
+                TaskInstance.queued_by_job_id,
+                TaskInstance.external_executor_id,
+            ).where(TaskInstance.id == workload_ti.id)
+        ).one_or_none()
+        if ti is None:
+            self.log.info(
+                "Dropping stale Kubernetes workload for %s because task instance id %s no longer exists",
+                task.key,
+                workload_ti.id,
+            )
+            return False
+
+        _, state, try_number, queued_by_job_id, external_executor_id = ti
+        if (
+            state == TaskInstanceState.QUEUED
+            and try_number == workload_ti.try_number
+            and queued_by_job_id == scheduler_job_id
+            and external_executor_id == launch_token
+        ):
+            return True
+
+        self.log.info(
+            "Dropping stale Kubernetes workload for %s because current task instance launch does not "
+            "match queued workload: ti_id=%s state=%s try_number=%s queued_by_job_id=%s "
+            "external_executor_id=%s workload_try_number=%s workload_external_executor_id=%s "
+            "scheduler_job_id=%s",
+            task.key,
+            workload_ti.id,
+            state,
+            try_number,
+            queued_by_job_id,
+            external_executor_id,
+            workload_ti.try_number,
+            launch_token,
+            scheduler_job_id,
+        )
+        return False
+
+    def _discard_stale_pod_creation_task(self, task: KubernetesJob) -> None:
+        """
+        Remove executor bookkeeping for a stale job that will not create a pod.
+
+        We intentionally do not fail the task instance here. Dropping the pod creation leaves the
+        row in its current DB state (typically still QUEUED under the *newer* launch that replaced
+        this workload). The newer launch owns the task from now on: either its own workload creates
+        the pod, or the scheduler's queued-task timeout / task-instance adoption reclaims the row.
+        Failing it here would clobber that newer launch.
+        """
+        self.running.discard(task.key)
+        queued_event = self.event_buffer.get(task.key)
+        if queued_event is not None and queued_event[0] == TaskInstanceState.QUEUED:
+            self.event_buffer.pop(task.key, None)
+        self.task_publish_retries.pop(task.key, None)
+        Stats.incr(
+            "kubernetes_executor.stale_workload_dropped",
+            tags=prune_dict({"team_name": self.team_name}),
+        )
+
     def sync(self) -> None:
         """Synchronize task state."""
         if TYPE_CHECKING:
@@ -426,6 +521,9 @@ def sync(self) -> None:
 
                 try:
                     key = task.key
+                    if not self._should_create_pod_for_job(task):
+                        self._discard_stale_pod_creation_task(task)
+                        continue
                     self.kube_scheduler.run_next(task)
                     self.task_publish_retries.pop(key, None)
                 except PodReconciliationError as e:

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py

@@ -575,11 +575,13 @@ def run_next(self, next_job: KubernetesJob) -> None:
         kube_image = next_job.kube_image or self.kube_config.kube_image
 
         dag_id, task_id, run_id, try_number, map_index = key
+        external_executor_id = None
         if len(command) == 1:
             from airflow.executors.workloads import ExecuteTask
 
             if isinstance(command[0], ExecuteTask):
                 workload = command[0]
+                external_executor_id = workload.ti.external_executor_id
                 command = workload_to_command_args(workload)
             else:
                 raise ValueError(
@@ -606,6 +608,7 @@ def run_next(self, next_job: KubernetesJob) -> None:
             map_index=map_index,
             date=None,
             run_id=run_id,
+            external_executor_id=external_executor_id,
             args=list(command),
             pod_override_object=kube_executor_config,
             base_worker_pod=base_worker_pod,

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/pod_generator.py

@@ -373,6 +373,7 @@ def construct_pod(
         scheduler_job_id: str,
         run_id: str | None = None,
         map_index: int = -1,
+        external_executor_id: str | None = None,
         *,
         with_mutation_hook: bool = False,
     ) -> k8s.V1Pod:
@@ -410,6 +411,8 @@ def construct_pod(
             annotations[get_logical_date_key()] = date.isoformat()
         if run_id:
             annotations["run_id"] = run_id
+        if external_executor_id:
+            annotations["external_executor_id"] = external_executor_id
 
         main_container = k8s.V1Container(
             name="base",