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 ad051b3e6d340..fb9ff13262943 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 @@ -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): 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 c1bac7960234d..3e6d2ccb3238a 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 @@ -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", diff --git a/airflow-core/src/airflow/executors/workloads/task.py b/airflow-core/src/airflow/executors/workloads/task.py index 68a917118e39c..6f82f9a8f68e8 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 542ce7eaaf15b..8e564cfad97e7 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 @@ -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 ): 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 46137915afbbe..a6f2b920397d8 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 @@ -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: 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 2c4bd306e52d7..0862b1f2799da 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 @@ -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, 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 285b263b064e1..a8a76d311f993 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 @@ -953,6 +953,90 @@ 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.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]) + 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 + 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" ) diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 74a680053711c..d9b3c3eb70c6f 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -247,9 +247,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 72aebeab3e49b..aa4294bed6a79 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 447f7c7594de6..1f5b1a3e1c0b9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1385,7 +1385,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 205b12bd4f2d8..acbb942ed7dea 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -344,6 +344,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"), @@ -351,7 +352,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 870b2e6deedca..1f9a69618b4e2 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -957,6 +957,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={ @@ -985,6 +987,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,