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..b3d716734123f 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 @@ -349,6 +349,85 @@ 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 an executor job still represents the current queued task instance. + + The scheduler creates an ``ExecuteTask`` workload while the task instance is queued, but the + Kubernetes pod may be created much later, for example after API-server throttling or quota + failures. In an HA scheduler deployment, the task instance may have been retried, cleared, or + otherwise replaced before this executor gets another chance to create the pod. Revalidating the + immutable task instance id and launch ownership here prevents an obsolete workload from creating + a stale worker pod. + """ + 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 + 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, + ).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 = ti + if ( + state == TaskInstanceState.QUEUED + and try_number == workload_ti.try_number + and queued_by_job_id == scheduler_job_id + ): + return True + + self.log.info( + "Dropping stale Kubernetes workload for %s because current task instance state does not " + "match the queued workload. task_instance_id=%s, state=%s, try_number=%s, " + "queued_by_job_id=%s, workload_try_number=%s, scheduler_job_id=%s", + task.key, + workload_ti.id, + state, + try_number, + queued_by_job_id, + workload_ti.try_number, + 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) + Stats.incr("kubernetes_executor.stale_workload_dropped") + def sync(self) -> None: """Synchronize task state.""" if TYPE_CHECKING: @@ -426,6 +505,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/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py index 285b263b064e1..7325dbb52aeb0 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 @@ -1029,6 +1029,49 @@ def test_skip_pod_creation_on_create_pods_after( 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_stale_execute_task_workload_before_pod_creation( + self, + mock_get_kube_client, + mock_kubernetes_job_watcher, + create_task_instance, + session, + ): + """A delayed Kubernetes workload should not create a pod after the DB task moved on.""" + 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 + session.merge(ti) + session.commit() + + workload = ExecuteTask.make(ti) + executor.queue_workload(workload, session=session) + executor._process_workloads([workload]) + + ti.state = TaskInstanceState.SUCCESS + 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" )