From 5e596ef7467ff2ebf1b09f9440fc942ae8f43c58 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Sun, 12 Jul 2026 19:44:57 +0200 Subject: [PATCH 01/16] Fence KubernetesExecutor launches with executor token --- .../execution_api/datamodels/taskinstance.py | 3 + .../execution_api/routes/task_instances.py | 15 ++++ .../src/airflow/executors/workloads/task.py | 2 +- .../versions/head/test_task_instances.py | 29 ++++++ .../executors/kubernetes_executor.py | 89 ++++++++++++++++++- .../executors/kubernetes_executor_utils.py | 3 + .../cncf/kubernetes/pod_generator.py | 3 + .../executors/test_kubernetes_executor.py | 43 +++++++++ task-sdk/src/airflow/sdk/api/client.py | 12 ++- .../airflow/sdk/api/datamodels/_generated.py | 2 + .../airflow/sdk/execution_time/supervisor.py | 7 +- task-sdk/tests/task_sdk/api/test_client.py | 3 +- .../execution_time/test_supervisor.py | 3 + 13 files changed, 207 insertions(+), 7 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index ddf31db971838..53fe7806bba14 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -65,6 +65,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 @@ -291,6 +293,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): 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 02d723a82185e..9aa531635f4f5 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 @@ -170,6 +170,7 @@ def ti_run( TI.unixname, TI.pid, TI.dag_version_id, + 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), @@ -197,6 +198,7 @@ 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) + payload_external_executor_id = data.pop("external_executor_id", None) # don't update start date when resuming from deferral if ti.next_kwargs: @@ -215,6 +217,19 @@ def ti_run( ti_run_payload.pid, ): log.info("Duplicate start request received", hostname=ti_run_payload.hostname) + elif 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", diff --git a/airflow-core/src/airflow/executors/workloads/task.py b/airflow-core/src/airflow/executors/workloads/task.py index 3099fe1d77485..7271a6e479b80 100644 --- a/airflow-core/src/airflow/executors/workloads/task.py +++ b/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? 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 80392fd17e206..b3da947c289c6 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 @@ -333,6 +333,35 @@ 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_returns_execution_token( self, client, exec_app, session, create_task_instance, time_machine ): diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index b75815df2fd91..b2f1a90908d89 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -37,7 +37,7 @@ from http import HTTPStatus from itertools import chain 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 @@ -104,6 +104,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. @@ -398,6 +399,82 @@ 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.""" + self.running.discard(task.key) + if self.event_buffer.get(task.key) == (TaskInstanceState.QUEUED, self.scheduler_job_id): + 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: @@ -486,6 +563,9 @@ def _create_pods_sequentially(self) -> None: task: KubernetesJob = self.task_queue.get_nowait() created += 1 try: + 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(task.key, None) except ( @@ -517,7 +597,12 @@ def _create_pods_concurrently(self) -> None: jobs: list[KubernetesJob] = [] with contextlib.suppress(Empty): for _ in range(self.kube_config.worker_pods_creation_batch_size): - jobs.append(self.task_queue.get_nowait()) + task = self.task_queue.get_nowait() + if not self._should_create_pod_for_job(task): + self._discard_stale_pod_creation_task(task) + self.task_queue.task_done() + continue + jobs.append(task) if not jobs: return start: float = time.monotonic() diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py index 0bf1ed8b23e6d..7df863d012acd 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py @@ -597,11 +597,13 @@ def _build_pod_request(self, next_job: KubernetesJob) -> k8s.V1Pod: 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( @@ -628,6 +630,7 @@ def _build_pod_request(self, next_job: KubernetesJob) -> k8s.V1Pod: 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, diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/pod_generator.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/pod_generator.py index ca499d363fc77..7f4458c99bd54 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/pod_generator.py +++ b/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", diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py index 5f1c03e7d4524..760eaa731508a 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py @@ -1012,6 +1012,49 @@ def test_run_next_exception_requeue( finally: kubernetes_executor.end() + @pytest.mark.db_test + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="workloads are used on Airflow 3+") + @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher") + @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") + def test_sync_drops_workload_with_stale_external_executor_id( + self, + mock_get_kube_client, + mock_kubernetes_job_watcher, + create_task_instance, + session, + ): + """A delayed Kubernetes workload must not create a pod after its launch token changes.""" + from airflow.executors.workloads import ExecuteTask + + executor = self.kubernetes_executor + executor.start() + try: + ti = create_task_instance(state=TaskInstanceState.QUEUED) + ti.queued_by_job_id = executor.job_id + ti.external_executor_id = "current-launch-token" + session.merge(ti) + session.commit() + + workload = ExecuteTask.make(ti) + executor.queue_workload(workload, session=session) + executor._process_workloads([workload]) + + ti.external_executor_id = "new-launch-token" + session.merge(ti) + session.commit() + + assert executor.kube_scheduler is not None + executor.kube_scheduler.run_next = mock.Mock() + executor.sync() + + executor.kube_scheduler.run_next.assert_not_called() + assert executor.task_queue is not None + assert executor.task_queue.empty() + assert ti.key not in executor.running + assert ti.key not in executor.event_buffer + finally: + executor.end() + @pytest.mark.skipif( AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed" ) diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index a0c0a30de088f..53b650851aca2 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -253,9 +253,17 @@ class TaskInstanceOperations: def __init__(self, client: Client): self.client = client - def start(self, id: uuid.UUID, pid: int, when: datetime) -> TIRunContext: + def start( + self, id: uuid.UUID, pid: int, when: datetime, external_executor_id: str | None = None + ) -> TIRunContext: """Tell the API server that this TI has started running.""" - body = TIEnterRunningPayload(pid=pid, hostname=get_hostname(), unixname=getuser(), start_date=when) + body = TIEnterRunningPayload( + pid=pid, + hostname=get_hostname(), + unixname=getuser(), + start_date=when, + external_executor_id=external_executor_id, + ) try: resp = self.client.patch(f"task-instances/{id}/run", content=body.model_dump_json()) diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 201f218c3c973..1ca936642efac 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -292,6 +292,7 @@ class TIEnterRunningPayload(BaseModel): unixname: Annotated[str, Field(title="Unixname")] pid: Annotated[int, Field(title="Pid")] start_date: Annotated[AwareDatetime, Field(title="Start Date")] + external_executor_id: Annotated[str | None, Field(title="External Executor Id")] = None class TIHeartbeatInfo(BaseModel): @@ -561,6 +562,7 @@ class TaskInstance(BaseModel): hostname: Annotated[str | None, Field(title="Hostname")] = None context_carrier: Annotated[dict[str, Any] | None, Field(title="Context Carrier")] = None queue: Annotated[str | None, Field(title="Queue")] = "default" + external_executor_id: Annotated[str | None, Field(title="External Executor Id")] = None class BundleInfo(BaseModel): diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 4c3ae38d21fab..07035806c10fa 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1448,7 +1448,12 @@ def _on_child_started( # We've forked, but the task won't start doing anything until we send it the StartupDetails # message. But before we do that, we need to tell the server it's started (so it has the chance to # tell us "no, stop!" for any reason) - ti_context = self.client.task_instances.start(ti.id, self.pid, datetime.now(tz=timezone.utc)) + ti_context = self.client.task_instances.start( + ti.id, + self.pid, + datetime.now(tz=timezone.utc), + getattr(ti, "external_executor_id", None), + ) self._should_retry = ti_context.should_retry self._last_successful_heartbeat = time.monotonic() except Exception: diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 8527349f56b26..7a2ccf724dc0c 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -361,6 +361,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert actual_body["pid"] == 100 assert actual_body["start_date"] == start_date assert actual_body["state"] == "running" + assert actual_body["external_executor_id"] == "launch-token" return httpx.Response( status_code=200, json=ti_context.model_dump(mode="json"), @@ -368,7 +369,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: return httpx.Response(status_code=400, json={"detail": "Bad Request"}) client = make_client(transport=httpx.MockTransport(handle_request)) - resp = client.task_instances.start(ti_id, 100, start_date) + resp = client.task_instances.start(ti_id, 100, start_date, external_executor_id="launch-token") assert resp == ti_context assert call_count == 3 diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index c9a977c9793c8..ec1f4ab6f0e93 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -959,6 +959,8 @@ def test_start_raises_task_already_running_and_kills_subprocess(self): def handle_request(request: httpx.Request) -> httpx.Response: if request.url.path == f"/task-instances/{ti_id}/run": + body = json.loads(request.read()) + assert body["external_executor_id"] == "launch-token" return httpx.Response( 409, json={ @@ -987,6 +989,7 @@ def subprocess_main(): try_number=1, dag_version_id=uuid7(), queue="default", + external_executor_id="launch-token", ), client=make_client(transport=httpx.MockTransport(handle_request)), target=subprocess_main, From dfdb876c566b5f8e9dbad9726b8cae050a37f7a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Mon, 13 Jul 2026 11:44:52 +0200 Subject: [PATCH 02/16] Make launch-token fence upgrade-safe and cover matching-token paths Only reject a task-instance /run when the worker actually presents an external_executor_id. Older Task SDK workers (mid rolling-upgrade, before the field existed) omit it entirely; fencing them would 409 every launch for pre-assigning executors (KubernetesExecutor, CeleryExecutor) until all workers are upgraded. Absent tokens now fall back to state-based validation. Also document that dropping a stale Kubernetes workload intentionally leaves the task instance in place (owned by the newer launch / reclaimed by queued-task timeout or adoption) rather than failing it. Tests: - API: matching token accepted; missing token accepted for old clients. - KubernetesExecutor: matching token creates the pod (positive counterpart to the stale-drop test). Co-Authored-By: Claude Opus 4.8 --- .../execution_api/routes/task_instances.py | 12 +++- .../versions/head/test_task_instances.py | 60 +++++++++++++++++++ .../executors/kubernetes_executor.py | 10 +++- .../executors/test_kubernetes_executor.py | 40 +++++++++++++ 4 files changed, 120 insertions(+), 2 deletions(-) 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 9aa531635f4f5..9b3b60e0a39a9 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 @@ -198,6 +198,12 @@ 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 @@ -217,7 +223,11 @@ def ti_run( ti_run_payload.pid, ): log.info("Duplicate start request received", hostname=ti_run_payload.hostname) - elif ti.external_executor_id is not None and ti.external_executor_id != payload_external_executor_id: + 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, 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 b3da947c289c6..dba2ead49108c 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 @@ -362,6 +362,66 @@ def test_ti_run_rejects_stale_external_executor_id(self, client, session, create 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 ): diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index b2f1a90908d89..987b48e674a26 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -465,7 +465,15 @@ def _should_create_pod_for_job(self, task: KubernetesJob, *, session: Session = 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.""" + """ + 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) if self.event_buffer.get(task.key) == (TaskInstanceState.QUEUED, self.scheduler_job_id): self.event_buffer.pop(task.key, None) diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py index 760eaa731508a..12c03b06496ca 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py @@ -1055,6 +1055,46 @@ def test_sync_drops_workload_with_stale_external_executor_id( finally: executor.end() + @pytest.mark.db_test + @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="workloads are used on Airflow 3+") + @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher") + @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") + def test_sync_creates_pod_when_external_executor_id_matches( + self, + mock_get_kube_client, + mock_kubernetes_job_watcher, + create_task_instance, + session, + ): + """A Kubernetes workload whose launch token still matches the DB must create a pod.""" + from airflow.executors.workloads import ExecuteTask + + executor = self.kubernetes_executor + executor.start() + try: + ti = create_task_instance(state=TaskInstanceState.QUEUED) + ti.queued_by_job_id = executor.job_id + ti.external_executor_id = "current-launch-token" + session.merge(ti) + session.commit() + + workload = ExecuteTask.make(ti) + executor.queue_workload(workload, session=session) + executor._process_workloads([workload]) + + # The DB token is left untouched, so the queued workload still owns the launch. + assert executor.kube_scheduler is not None + executor.kube_scheduler.run_next = mock.Mock() + executor.sync() + + executor.kube_scheduler.run_next.assert_called_once() + created_job = executor.kube_scheduler.run_next.call_args.args[0] + assert created_job.key == ti.key + assert executor.task_queue is not None + assert executor.task_queue.empty() + finally: + executor.end() + @pytest.mark.skipif( AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed" ) From 42fa8976c1a5677d1bd91bf4626d9ad67e423c5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Mon, 13 Jul 2026 14:12:43 +0200 Subject: [PATCH 03/16] Preserve KubernetesExecutor launch token in queued events --- .../kubernetes/executors/kubernetes_executor.py | 14 ++++++++++++-- .../executors/test_kubernetes_executor.py | 1 + 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index 987b48e674a26..262eccb5c2f42 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -368,9 +368,18 @@ def execute_async( queue, ) - self.event_buffer[key] = (TaskInstanceState.QUEUED, self.scheduler_job_id) job = KubernetesJob(key, command, kube_executor_config, pod_template_file, coordinator_kube_image) self.pod_launch_attempts[key] = _PodLaunchAttempt(job=job) + 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 (ImportError, TypeError): + pass + + self.event_buffer[key] = (TaskInstanceState.QUEUED, event_info) self.task_queue.put(job) def queue_workload(self, workload: workloads.All, session: Session | None) -> None: @@ -475,7 +484,8 @@ def _discard_stale_pod_creation_task(self, task: KubernetesJob) -> None: Failing it here would clobber that newer launch. """ self.running.discard(task.key) - if self.event_buffer.get(task.key) == (TaskInstanceState.QUEUED, self.scheduler_job_id): + 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( diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py index 12c03b06496ca..0eff5625ffe69 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py @@ -1081,6 +1081,7 @@ def test_sync_creates_pod_when_external_executor_id_matches( workload = ExecuteTask.make(ti) executor.queue_workload(workload, session=session) executor._process_workloads([workload]) + assert executor.event_buffer[ti.key] == (TaskInstanceState.QUEUED, "current-launch-token") # The DB token is left untouched, so the queued workload still owns the launch. assert executor.kube_scheduler is not None From 2aa8f22c80b504f40ac8e3a60d24e727c821d4ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Mon, 20 Jul 2026 12:59:53 +0200 Subject: [PATCH 04/16] Fix launch-token CI failures --- .../kubernetes/executors/kubernetes_executor.py | 17 ++++++++++++----- .../edge3/worker_api/v2-edge-generated.yaml | 5 +++++ .../observability/metrics/metrics_template.yaml | 6 ++++++ .../sdk/execution_time/schema/schema.json | 12 ++++++++++++ .../task_sdk/execution_time/test_supervisor.py | 2 +- 5 files changed, 36 insertions(+), 6 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index 262eccb5c2f42..9f58905793f5c 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -373,11 +373,14 @@ def execute_async( 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 (ImportError, TypeError): + except ImportError: pass + else: + try: + 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(job) @@ -411,7 +414,11 @@ def _process_workloads(self, workloads: Sequence[workloads.All]) -> None: @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 + try: + from airflow.executors.workloads import ExecuteTask + except ImportError: + return True + from airflow.models.taskinstance import TaskInstance if not task.command or not isinstance(task.command[0], ExecuteTask): diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml index 13cacba97be36..8f7c05aa6a6b0 100644 --- a/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml +++ b/providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml @@ -1318,6 +1318,11 @@ components: type: string title: Queue default: default + external_executor_id: + anyOf: + - type: string + - type: 'null' + title: External Executor Id pool_slots: type: integer title: Pool Slots diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml index 392b55d7c072f..5d7dc53f32cc5 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -641,6 +641,12 @@ metrics: legacy_name: "-" name_variables: ["status"] + - name: "kubernetes_executor.stale_workload_dropped" + description: "Number of stale queued KubernetesExecutor workloads dropped before worker pod creation." + type: "counter" + legacy_name: "-" + name_variables: [] + # ========== # Timers # ========== diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 4524c74ff794a..d48ad9a916ee2 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -5176,6 +5176,18 @@ "default": "default", "title": "Queue", "type": "string" + }, + "external_executor_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "External Executor Id" } }, "required": [ diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index ec1f4ab6f0e93..1a78c63379519 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -803,7 +803,7 @@ def mock_monotonic(): assert exit_code == 0, captured_logs # Validate calls to the client - mock_client.task_instances.start.assert_called_once_with(ti.id, mocker.ANY, mocker.ANY) + mock_client.task_instances.start.assert_called_once_with(ti.id, mocker.ANY, mocker.ANY, None) mock_client.task_instances.heartbeat.assert_called_once_with(ti.id, pid=mocker.ANY) mock_client.task_instances.defer.assert_called_once_with( ti.id, From 0f46519c7466c7751486423fce5382763382f930 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Mon, 20 Jul 2026 16:09:55 +0200 Subject: [PATCH 05/16] Add execution API version for launch token --- .../api_fastapi/execution_api/versions/__init__.py | 2 ++ .../execution_api/versions/v2026_06_30.py | 12 ++++++++++++ 2 files changed, 14 insertions(+) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index d56ec735c8f13..c72d38d84078b 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -47,6 +47,7 @@ AddPartitionDateField, AddRetryPolicyFields, AddTaskAndAssetStateStoreEndpoints, + AddTaskInstanceExternalExecutorIdField, AddTaskInstanceQueueField, AddTeamNameField, AddVariableKeysEndpoint, @@ -61,6 +62,7 @@ AddVariableKeysEndpoint, AddConnectionTestEndpoint, AddAwaitingInputStatePayload, + AddTaskInstanceExternalExecutorIdField, AddTaskInstanceQueueField, AddRetryPolicyFields, AddTeamNameField, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py index 0a31681069161..b991fb13f24ba 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py @@ -29,6 +29,7 @@ DagRun, TaskInstance, TIAwaitingInputStatePayload, + TIEnterRunningPayload, TIRetryStatePayload, TIRunContext, ) @@ -61,6 +62,17 @@ class AddTaskInstanceQueueField(VersionChange): instructions_to_migrate_to_previous_version = (schema(TaskInstance).field("queue").didnt_exist,) +class AddTaskInstanceExternalExecutorIdField(VersionChange): + """Add the `external_executor_id` field to task instance launch payloads.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(TaskInstance).field("external_executor_id").didnt_exist, + schema(TIEnterRunningPayload).field("external_executor_id").didnt_exist, + ) + + class AddAwaitingInputStatePayload(VersionChange): """Add the awaiting_input task instance state transition payload (Human-in-the-loop, no trigger).""" From 8b5abeb492103050d469bf966700103767c9b66b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Mon, 20 Jul 2026 16:51:31 +0200 Subject: [PATCH 06/16] Add supervisor schema version marker --- .../schema/versions/v2026_06_16.py | 25 +++++++++++++++++++ ts-sdk/src/generated/supervisor.ts | 2 ++ 2 files changed, 27 insertions(+) create mode 100644 task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_06_16.py diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_06_16.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_06_16.py new file mode 100644 index 0000000000000..c62fd6551f537 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_06_16.py @@ -0,0 +1,25 @@ +# 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. + +""" +Supervisor schema changes for the initial 2026-06-16 schema version. + +Cadwyn does not allow VersionChange entries on the first supervisor schema +version because there is no previous version to migrate to. This file is kept +with the schema-changing PR so the supervisor schema version hook records that +the current head-version schema was intentionally updated. +""" diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index e12e5d815ae20..c987e5ff7ca26 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -195,6 +195,7 @@ export type ContextCarrier = { [k: string]: unknown; } | null; export type Queue = string; +export type ExternalExecutorId = string | null; export type IsFailureCallback = boolean | null; export type Type12 = "DagCallbackRequest"; export type File = string; @@ -977,6 +978,7 @@ export interface TaskInstance { hostname?: Hostname; context_carrier?: ContextCarrier; queue?: Queue; + external_executor_id?: ExternalExecutorId; } /** * Request for DAG File Parsing. From 4b3b461c339ba00e372f34f8e9d698c8b3c466bb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Mon, 20 Jul 2026 17:29:54 +0200 Subject: [PATCH 07/16] Update workload serialization test for launch token --- airflow-core/tests/unit/executors/test_workloads.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow-core/tests/unit/executors/test_workloads.py b/airflow-core/tests/unit/executors/test_workloads.py index 37fbcd96ce950..6d41a472dd2a2 100644 --- a/airflow-core/tests/unit/executors/test_workloads.py +++ b/airflow-core/tests/unit/executors/test_workloads.py @@ -159,7 +159,7 @@ def test_workload_ti_round_trips_through_sdk_generated_model(): ) dumped = ti.model_dump(mode="json") - assert "external_executor_id" not in dumped + assert dumped["external_executor_id"] == "celery-id" assert "executor_config" not in dumped # Executor-side scheduling fields stay on the workload wire (older workers # deserialize the workload with a model that requires them) but are not @@ -169,6 +169,7 @@ def test_workload_ti_round_trips_through_sdk_generated_model(): received = GeneratedTaskInstance.model_validate(dumped) assert received.queue == "jdk-17" + assert received.external_executor_id == "celery-id" assert received.map_index == 3 assert not hasattr(received, "pool_slots") From b78bbb5e6c805b17571ae4d4b92f1dec387e0183 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Tue, 21 Jul 2026 00:09:17 +0200 Subject: [PATCH 08/16] Handle launch tokens on older Airflow workloads --- .../cncf/kubernetes/executors/kubernetes_executor.py | 6 +++--- .../cncf/kubernetes/executors/test_kubernetes_executor.py | 8 ++++++-- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index 9f58905793f5c..c737b553f46be 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -378,7 +378,7 @@ def execute_async( else: try: if len(command) == 1 and isinstance(command[0], ExecuteTask): - event_info = command[0].ti.external_executor_id or self.scheduler_job_id + event_info = getattr(command[0].ti, "external_executor_id", None) or self.scheduler_job_id except TypeError: pass @@ -426,7 +426,7 @@ def _should_create_pod_for_job(self, task: KubernetesJob, *, session: Session = workload = task.command[0] workload_ti = workload.ti - launch_token = workload_ti.external_executor_id + launch_token = getattr(workload_ti, "external_executor_id", None) try: scheduler_job_id = int(self.scheduler_job_id) if self.scheduler_job_id is not None else None @@ -459,7 +459,7 @@ def _should_create_pod_for_job(self, task: KubernetesJob, *, session: Session = state == TaskInstanceState.QUEUED and try_number == workload_ti.try_number and queued_by_job_id == scheduler_job_id - and external_executor_id == launch_token + and (launch_token is None or external_executor_id == launch_token) ): return True diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py index 0eff5625ffe69..2cfa0553ef0ae 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py @@ -1013,7 +1013,9 @@ def test_run_next_exception_requeue( kubernetes_executor.end() @pytest.mark.db_test - @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="workloads are used on Airflow 3+") + @pytest.mark.skipif( + not AIRFLOW_V_3_2_PLUS, reason="workloads include external_executor_id on Airflow 3.2+" + ) @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher") @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") def test_sync_drops_workload_with_stale_external_executor_id( @@ -1056,7 +1058,9 @@ def test_sync_drops_workload_with_stale_external_executor_id( executor.end() @pytest.mark.db_test - @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="workloads are used on Airflow 3+") + @pytest.mark.skipif( + not AIRFLOW_V_3_2_PLUS, reason="workloads include external_executor_id on Airflow 3.2+" + ) @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher") @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") def test_sync_creates_pod_when_external_executor_id_matches( From 377426a6d075be1c706e2d8216252e5df51b3bef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Tue, 21 Jul 2026 09:51:20 +0200 Subject: [PATCH 09/16] Avoid DB access for legacy Kubernetes jobs --- .../cncf/kubernetes/executors/kubernetes_executor.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index c737b553f46be..bfa5e77004880 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -63,7 +63,7 @@ from airflow.providers.common.compat.sdk import Stats, conf from airflow.utils.helpers import prune_dict from airflow.utils.log.logging_mixin import remove_escape_codes -from airflow.utils.session import NEW_SESSION, provide_session +from airflow.utils.session import NEW_SESSION, create_session, provide_session from airflow.utils.state import TaskInstanceState if TYPE_CHECKING: @@ -411,8 +411,7 @@ 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: + def _should_create_pod_for_job(self, task: KubernetesJob, *, session: Session | None = None) -> bool: """Check whether a queued Kubernetes job still owns the current task launch token.""" try: from airflow.executors.workloads import ExecuteTask @@ -428,6 +427,10 @@ def _should_create_pod_for_job(self, task: KubernetesJob, *, session: Session = workload_ti = workload.ti launch_token = getattr(workload_ti, "external_executor_id", None) + if session is None: + with create_session() as session: + return self._should_create_pod_for_job(task, session=session) + try: scheduler_job_id = int(self.scheduler_job_id) if self.scheduler_job_id is not None else None except ValueError: From 878e218469a42e15e9ea37ddf1dfe09769317e6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Fri, 24 Jul 2026 19:14:56 +0200 Subject: [PATCH 10/16] Remove stale workload metric from launch-token PR --- .../cncf/kubernetes/executors/kubernetes_executor.py | 4 ---- .../observability/metrics/metrics_template.yaml | 6 ------ 2 files changed, 10 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index bfa5e77004880..a90c0fbe9924a 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -498,10 +498,6 @@ def _discard_stale_pod_creation_task(self, task: KubernetesJob) -> None: 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.""" diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml index 5d7dc53f32cc5..392b55d7c072f 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -641,12 +641,6 @@ metrics: legacy_name: "-" name_variables: ["status"] - - name: "kubernetes_executor.stale_workload_dropped" - description: "Number of stale queued KubernetesExecutor workloads dropped before worker pod creation." - type: "counter" - legacy_name: "-" - name_variables: [] - # ========== # Timers # ========== From b99ea3be1507fbfdffbac8a235b33bfc9f2c3970 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Tue, 28 Jul 2026 21:18:27 +0200 Subject: [PATCH 11/16] Allow null launch token during rolling upgrades --- .../execution_api/routes/task_instances.py | 14 ++++----- .../versions/head/test_task_instances.py | 30 +++++++++++++++++++ 2 files changed, 37 insertions(+), 7 deletions(-) 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 9b3b60e0a39a9..10f7074d77735 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 @@ -198,12 +198,12 @@ 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 + # Only fence on the launch token when the worker actually presented a non-null token. Older + # Task SDK clients may omit the field entirely, and newer clients can serialize it as null + # before scheduler-side preassignment is available everywhere. Treating either case as a stale + # launch would 409 every task start for pre-assigning executors (KubernetesExecutor, + # CeleryExecutor) during rolling upgrades. When it is absent/null we fall back to state-based + # validation, matching the pre-token behavior. payload_external_executor_id = data.pop("external_executor_id", None) # don't update start date when resuming from deferral @@ -224,7 +224,7 @@ def ti_run( ): log.info("Duplicate start request received", hostname=ti_run_payload.hostname) elif ( - payload_has_external_executor_id + payload_external_executor_id is not None and ti.external_executor_id is not None and ti.external_executor_id != payload_external_executor_id ): 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 dba2ead49108c..fdaf09b22d649 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 @@ -422,6 +422,36 @@ def test_ti_run_allows_missing_external_executor_id_for_old_clients( session.refresh(ti) assert ti.state == State.RUNNING + def test_ti_run_allows_null_external_executor_id_for_rolling_upgrade( + self, client, session, create_task_instance + ): + """A worker sending a null launch token must not be fenced out during rolling upgrades.""" + ti = create_task_instance( + task_id="test_ti_run_allows_null_external_executor_id_for_rolling_upgrade", + 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": None, + }, + ) + + 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 ): From 803aa4989b440a8db4d8f58f118aff69c257e8ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Tue, 28 Jul 2026 22:58:27 +0200 Subject: [PATCH 12/16] Prevent TinkerPop provider downgrade in constraints --- .../in_container/run_generate_constraints.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/scripts/in_container/run_generate_constraints.py b/scripts/in_container/run_generate_constraints.py index 6bd75777253f4..7e79c9ab0d469 100755 --- a/scripts/in_container/run_generate_constraints.py +++ b/scripts/in_container/run_generate_constraints.py @@ -620,14 +620,30 @@ def generate_constraints_pypi_providers(config_params: ConfigParams) -> None: # ``opentelemetry-*>=1.27.0`` floor — opentelemetry-sdk exact-pins the semantic conventions # version it ships with, so a higher floor here would conflict with any sdk older than that pairing. # - # These two are the only pre-releases in the constraints we tag, and removing the need for the - # exception is tracked at https://github.com/apache/airflow/issues/71176 - # + # Removing the need for the opentelemetry pre-release exception is tracked at + # https://github.com/apache/airflow/issues/71176 + # * pymysql<1.2 — pymysql 1.2.0 changed Connection.ping() to require `reconnect` as a + # positional arg, which breaks SQLAlchemy's AsyncAdapt_aiomysql_connection.ping() (it + # has no default for `reconnect`). The released apache-airflow-providers-mysql on PyPI + # does not yet carry this cap, so we mirror it here so PyPI constraints stay installable + # until the SQLAlchemy fix is released. Tracked upstream at + # https://github.com/sqlalchemy/sqlalchemy/issues/13306 + # * databricks-sql-connector>=4.0.0 - added to keep databricks-sql-connector from downgrading + # because of https://github.com/apache/thrift/pull/3584 - which shipped thrift 0.24.0. Older + # versions of databricks-sql-connector do not have the thrift<=0.23.0 limitation, so the + # resolver preferred the latest thrift over the latest connector and downgraded the connector. + # This is tracked in https://github.com/databricks/databricks-sql-python/issues/859 + # * apache-airflow-providers-apache-tinkerpop>=1.1.4 - added to prevent the highest-resolution + # constraints job from downgrading the released TinkerPop provider to 1.0.2 while resolving the + # latest PyPI provider set. additional_constraints_for_highest_resolution: list[str] = [ "pyarrow>=22.0.0; python_version >= '3.14'", "gremlinpython>=3.8.0", "opentelemetry-exporter-prometheus>=0.47b0", "opentelemetry-semantic-conventions>=0.48b0", + "pymysql>=1.0.3,<1.2", + "databricks-sql-connector>=4.0.0", + "apache-airflow-providers-apache-tinkerpop>=1.1.4", ] # Every run pins the providers at the versions PyPI holds when it is generated. Only a run for From 5926ca83661da79f39be66eea783c6780010f239 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Thu, 30 Jul 2026 11:12:18 +0200 Subject: [PATCH 13/16] Keep KubernetesExecutor compatible with Airflow 3.0 and 3.1 --- .../executors/kubernetes_executor.py | 33 ++++++++++--------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index a90c0fbe9924a..55a6879c9d172 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -59,7 +59,7 @@ annotations_to_key, ) from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator -from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_0_PLUS +from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS from airflow.providers.common.compat.sdk import Stats, conf from airflow.utils.helpers import prune_dict from airflow.utils.log.logging_mixin import remove_escape_codes @@ -440,15 +440,15 @@ def _should_create_pod_for_job(self, task: KubernetesJob, *, session: Session | ) 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() + columns = [ + TaskInstance.id, + TaskInstance.state, + TaskInstance.try_number, + TaskInstance.queued_by_job_id, + ] + if AIRFLOW_V_3_2_PLUS: + columns.append(TaskInstance.external_executor_id) + ti = session.execute(select(*columns).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", @@ -457,27 +457,28 @@ def _should_create_pod_for_job(self, task: KubernetesJob, *, session: Session | ) return False - _, state, try_number, queued_by_job_id, external_executor_id = ti + _, state, try_number, queued_by_job_id, *optional_values = ti + external_executor_id = optional_values[0] if optional_values else None if ( state == TaskInstanceState.QUEUED and try_number == workload_ti.try_number and queued_by_job_id == scheduler_job_id - and (launch_token is None or external_executor_id == launch_token) + and (not AIRFLOW_V_3_2_PLUS or launch_token is None or 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", + "match queued workload: " + "ti_id=%s state=%s try_number=%s queued_by_job_id=%s workload_try_number=%s " + "external_executor_id=%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, + external_executor_id, launch_token, scheduler_job_id, ) From 50fd2086012dad1e1ba80aa18183a9efbf0c1635 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Thu, 30 Jul 2026 13:58:16 +0200 Subject: [PATCH 14/16] Detect launch-token support from the installed model --- .../cncf/kubernetes/executors/kubernetes_executor.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index 55a6879c9d172..9bf2cf1fbca15 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -59,7 +59,7 @@ annotations_to_key, ) from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator -from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS +from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_0_PLUS from airflow.providers.common.compat.sdk import Stats, conf from airflow.utils.helpers import prune_dict from airflow.utils.log.logging_mixin import remove_escape_codes @@ -440,13 +440,14 @@ def _should_create_pod_for_job(self, task: KubernetesJob, *, session: Session | ) return True + has_external_executor_id = hasattr(TaskInstance, "external_executor_id") columns = [ TaskInstance.id, TaskInstance.state, TaskInstance.try_number, TaskInstance.queued_by_job_id, ] - if AIRFLOW_V_3_2_PLUS: + if has_external_executor_id: columns.append(TaskInstance.external_executor_id) ti = session.execute(select(*columns).where(TaskInstance.id == workload_ti.id)).one_or_none() if ti is None: @@ -463,7 +464,7 @@ def _should_create_pod_for_job(self, task: KubernetesJob, *, session: Session | state == TaskInstanceState.QUEUED and try_number == workload_ti.try_number and queued_by_job_id == scheduler_job_id - and (not AIRFLOW_V_3_2_PLUS or launch_token is None or external_executor_id == launch_token) + and (not has_external_executor_id or launch_token is None or external_executor_id == launch_token) ): return True From 0e1f98189c5464448ad6b5d94136c278d25e2190 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Thu, 30 Jul 2026 22:35:00 +0200 Subject: [PATCH 15/16] Exit cleanly when a task instance is superseded before startup --- task-sdk/src/airflow/sdk/api/client.py | 13 ++++- task-sdk/src/airflow/sdk/exceptions.py | 12 +++++ .../airflow/sdk/execution_time/supervisor.py | 15 +++++- task-sdk/tests/task_sdk/api/test_client.py | 41 ++++++++++++++- .../execution_time/test_supervisor.py | 52 ++++++++++++++++++- 5 files changed, 128 insertions(+), 5 deletions(-) diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 53b650851aca2..b5274b2fdd40b 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -90,7 +90,7 @@ XComSequenceSliceResponse, ) from airflow.sdk.configuration import conf -from airflow.sdk.exceptions import ErrorType, TaskAlreadyRunningError +from airflow.sdk.exceptions import ErrorType, TaskAlreadyRunningError, TaskInstanceSupersededError from airflow.sdk.execution_time.comms import ( AssetsByAliasResult, CreateHITLDetailPayload, @@ -185,6 +185,7 @@ def getuser() -> str: "ConnectionOperations", "ServerResponseError", "TaskInstanceOperations", + "TaskInstanceSupersededError", "get_hostname", "getuser", ] @@ -268,14 +269,22 @@ def start( try: resp = self.client.patch(f"task-instances/{id}/run", content=body.model_dump_json()) except ServerResponseError as e: + detail = e.detail if e.response.status_code == HTTPStatus.CONFLICT: - detail = e.detail if ( isinstance(detail, dict) and detail.get("reason") == "invalid_state" and detail.get("previous_state") == "running" ): raise TaskAlreadyRunningError(f"Task instance {id} is already running") from e + if isinstance(detail, dict) and detail.get("reason") == "stale_executor_launch": + raise TaskInstanceSupersededError( + f"Task instance {id} was superseded by a newer launch" + ) from e + if e.response.status_code == HTTPStatus.NOT_FOUND and ( + not isinstance(detail, dict) or detail.get("reason") == "not_found" + ): + raise TaskInstanceSupersededError(f"Task instance {id} no longer exists") from e raise return TIRunContext.model_validate_json(resp.read()) diff --git a/task-sdk/src/airflow/sdk/exceptions.py b/task-sdk/src/airflow/sdk/exceptions.py index 6f43d5421ecf2..951168d715f0a 100644 --- a/task-sdk/src/airflow/sdk/exceptions.py +++ b/task-sdk/src/airflow/sdk/exceptions.py @@ -403,6 +403,18 @@ class TaskAlreadyRunningError(AirflowException): """Raised when a task is already running on another worker.""" +class TaskInstanceSupersededError(AirflowException): + """ + Raised when a worker's task instance launch has been superseded before it started running. + + This happens when the task instance the worker was launched for no longer exists + (``/run`` returns ``404 not_found`` because a newer attempt replaced it) or the worker + no longer owns the current launch (``/run`` returns ``409 stale_executor_launch``). In + both cases the worker never executed the task, so it should exit cleanly without + consuming a task retry. + """ + + class FailFastDagInvalidTriggerRule(AirflowException): """Raise when a dag has 'fail_fast' enabled yet has a non-default trigger rule.""" diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 07035806c10fa..328f31877e9c5 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -58,7 +58,7 @@ TaskInstanceState, ) from airflow.sdk.configuration import conf -from airflow.sdk.exceptions import ErrorType +from airflow.sdk.exceptions import ErrorType, TaskInstanceSupersededError from airflow.sdk.execution_time import comms from airflow.sdk.execution_time.comms import ( AssetEventsResult, @@ -2636,6 +2636,19 @@ def supervise_task( final_state=result.final_state, ) return result.exit_code + except TaskInstanceSupersededError: + # The task instance this worker was launched for was superseded before it started + # running (``/run`` returned 404 not_found or 409 stale_executor_launch). The task + # never executed, so exit cleanly (code 0) to avoid surfacing a pod failure and + # consuming a task retry. Returning 0 here (rather than propagating) keeps the pod in + # a Succeeded/Completed state; the stale event carries the old try_number and is + # ignored by the scheduler, so the live attempt is unaffected. + log.info( + "Task instance superseded before startup; exiting without consuming a retry", + workload_type="ExecuteTask", + workload_id=str(ti.id), + ) + return 0 finally: if log_path and log_file_descriptor: log_file_descriptor.close() diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 7a2ccf724dc0c..6c9a0ec226fd1 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -50,7 +50,7 @@ VariableResponse, XComResponse, ) -from airflow.sdk.exceptions import ErrorType, TaskAlreadyRunningError +from airflow.sdk.exceptions import ErrorType, TaskAlreadyRunningError, TaskInstanceSupersededError from airflow.sdk.execution_time.comms import ( AssetsByAliasResult, DeferTask, @@ -420,6 +420,45 @@ def handle_request(request: httpx.Request) -> httpx.Response: with pytest.raises(ServerResponseError): client.task_instances.start(ti_id, 100, datetime(2024, 10, 31, tzinfo=timezone.utc)) + def test_task_instance_start_not_found(self): + """Test that start() raises TaskInstanceSupersededError on a 404 not_found response.""" + ti_id = uuid6.uuid7() + + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == f"/task-instances/{ti_id}/run": + return httpx.Response( + 404, + json={"detail": {"reason": "not_found", "message": "Task Instance not found"}}, + ) + return httpx.Response(status_code=204) + + client = make_client(transport=httpx.MockTransport(handle_request)) + + with pytest.raises(TaskInstanceSupersededError, match="no longer exists"): + client.task_instances.start(ti_id, 100, datetime(2024, 10, 31, tzinfo=timezone.utc)) + + def test_task_instance_start_stale_executor_launch(self): + """Test that start() raises TaskInstanceSupersededError on a 409 stale_executor_launch response.""" + ti_id = uuid6.uuid7() + + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == f"/task-instances/{ti_id}/run": + return httpx.Response( + 409, + json={ + "detail": { + "reason": "stale_executor_launch", + "message": "TI executor launch token does not match the current queued task instance", + } + }, + ) + return httpx.Response(status_code=204) + + client = make_client(transport=httpx.MockTransport(handle_request)) + + with pytest.raises(TaskInstanceSupersededError, match="superseded"): + client.task_instances.start(ti_id, 100, datetime(2024, 10, 31, tzinfo=timezone.utc)) + @pytest.mark.parametrize( "state", [state for state in TerminalTIState if state != TerminalTIState.SUCCESS] ) diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 1a78c63379519..69e23e34271b8 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -68,7 +68,12 @@ TaskInstance, TaskInstanceState, ) -from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType, TaskAlreadyRunningError +from airflow.sdk.exceptions import ( + AirflowRuntimeError, + ErrorType, + TaskAlreadyRunningError, + TaskInstanceSupersededError, +) from airflow.sdk.execution_time import supervisor, task_runner from airflow.sdk.execution_time.comms import ( AssetEventsResult, @@ -995,6 +1000,51 @@ def subprocess_main(): target=subprocess_main, ) + @pytest.mark.parametrize( + ("status_code", "detail"), + [ + pytest.param(404, {"reason": "not_found", "message": "Task Instance not found"}, id="not_found"), + pytest.param( + 409, + { + "reason": "stale_executor_launch", + "message": "TI executor launch token does not match the current queued task instance", + }, + id="stale_executor_launch", + ), + ], + ) + def test_start_raises_task_instance_superseded_and_kills_subprocess(self, status_code, detail): + """Test that ActivitySubprocess.start() raises TaskInstanceSupersededError and kills the child + when the API reports the task instance was superseded (404 not_found or 409 stale_executor_launch).""" + ti_id = uuid7() + + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == f"/task-instances/{ti_id}/run": + return httpx.Response(status_code, json={"detail": detail}) + return httpx.Response(status_code=204) + + def subprocess_main(): + CommsDecoder()._get_response() + + with pytest.raises(TaskInstanceSupersededError): + ActivitySubprocess.start( + dag_rel_path=os.devnull, + bundle_info=FAKE_BUNDLE, + what=TaskInstance( + id=ti_id, + task_id="b", + dag_id="c", + run_id="d", + try_number=1, + dag_version_id=uuid7(), + queue="default", + external_executor_id="launch-token", + ), + client=make_client(transport=httpx.MockTransport(handle_request)), + target=subprocess_main, + ) + @pytest.mark.parametrize("captured_logs", [logging.WARNING], indirect=True) def test_heartbeat_failures_handling(self, monkeypatch, mocker, captured_logs, time_machine): """ From ca4b537509e6e0b1ffddb1c95ddaf3f413509622 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Guilherme=20Da=20Silva=20Gon=C3=A7alves?= Date: Fri, 31 Jul 2026 14:27:21 +0200 Subject: [PATCH 16/16] Keep launch-token pod annotation compatible with older Airflow --- .../cncf/kubernetes/executors/kubernetes_executor_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py index 7df863d012acd..0c14047781b7f 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py @@ -603,7 +603,10 @@ def _build_pod_request(self, next_job: KubernetesJob) -> k8s.V1Pod: if isinstance(command[0], ExecuteTask): workload = command[0] - external_executor_id = workload.ti.external_executor_id + # ``external_executor_id`` is only present on the workload TI in Airflow 3.2+; on + # older versions exercised by provider compatibility jobs it is absent, so fall back + # to ``None`` (the launch-token annotation is simply omitted from the pod there). + external_executor_id = getattr(workload.ti, "external_executor_id", None) command = workload_to_command_args(workload) else: raise ValueError(