eladkal commented on code in PR #69762:
URL: https://github.com/apache/airflow/pull/69762#discussion_r3628663846
##########
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py:
##########
@@ -391,6 +391,94 @@ 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)
+ def _should_create_pod_for_job(self, task: KubernetesJob) -> 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.
+ """
+ try:
+ from airflow.executors.workloads import ExecuteTask
+ except ImportError:
+ # Compatibility with older Airflow versions tested by provider
compatibility jobs.
+ return True
+
+ if not task.command or not isinstance(task.command[0], ExecuteTask):
+ return True
+
+ return self._should_create_pod_for_execute_task(task, task.command[0])
+
+ @provide_session
+ def _should_create_pod_for_execute_task(
+ self,
+ task: KubernetesJob,
+ workload: Any,
+ *,
+ session: Session = NEW_SESSION,
+ ) -> bool:
+ """Check that an ``ExecuteTask`` workload still owns the queued task
instance row."""
+ from airflow.models.taskinstance import TaskInstance
+
+ 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(cast(TaskInstance.id, String) == str(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")
Review Comment:
I guess this is not backward compatible as it will be depended on newer
Airflow version that ships this metric?
##########
shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml:
##########
@@ -632,6 +632,12 @@ metrics:
legacy_name: "-"
name_variables: ["status"]
+ - name: "kubernetes_executor.stale_workload_dropped"
+ description: "Number of stale KubernetesExecutor workloads dropped before
worker pod creation."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
Review Comment:
We try to avoid having changes to providers and airflow core parts on the
same PR as they don't have the same release cycle. I suggest to extract this
into a separated PR. When this part is merged and schedule for release then the
provider code can be adjusted under a version check
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]