potiuk commented on code in PR #69762:
URL: https://github.com/apache/airflow/pull/69762#discussion_r3683273191
##########
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py:
##########
@@ -398,6 +398,93 @@ 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))
Review Comment:
This needs to change before merge. `TaskInstance.id` is `Uuid()` and the
table's **primary key**:
```python
id: Mapped[UUID] = mapped_column(Uuid(), primary_key=True, default=uuid7)
```
Wrapping it in `cast(..., String)` makes the predicate non-sargable — the
database can no longer use the primary key index and must cast every row in
`task_instance` to text to compare. That is a **full scan of the largest table
in most deployments, on every pod creation**, which is a worse operational
problem than the bug being fixed.
Compare the UUID directly and coerce the *parameter* instead of the column:
```python
).where(TaskInstance.id == workload_ti.id)
```
If `workload_ti.id` arrives as a `str` rather than a `UUID`, normalise it
once (`UUID(str(workload_ti.id))`) before the query — never by casting the
column.
Worth adding a brief note to the PR description about the added per-pod
query as well: even once this is indexed, it's one extra round-trip per pod
launch, which is a reasonable price for correctness but should be a conscious
one.
---
Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting
##########
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py:
##########
@@ -398,6 +398,93 @@ 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,
Review Comment:
`Any` here loses the type information the function depends on — it reads
`workload.ti.id` and `workload.ti.try_number`, so a change to `ExecuteTask`'s
shape would fail silently at runtime rather than in mypy.
Since `ExecuteTask` is already imported lazily in the caller for
version-compat reasons, the usual pattern is a `TYPE_CHECKING` import plus a
string annotation:
```python
if TYPE_CHECKING:
from airflow.executors.workloads import ExecuteTask
def _should_create_pod_for_execute_task(self, task: KubernetesJob, workload:
ExecuteTask, ...)
```
Keeps the runtime compat shim and gets the checking back.
---
Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting
--
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]
