You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2026/07/29 09:58:57 UTC
(airflow) branch main updated: Add opt-in concurrent pod creation to KubernetesExecutor (#68480)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 703a0570c84 Add opt-in concurrent pod creation to KubernetesExecutor (#68480)
703a0570c84 is described below
commit 703a0570c845df8d99122d643ee9cf600562fe54
Author: Stefan Wang <1f...@gmail.com>
AuthorDate: Wed Jul 29 02:58:46 2026 -0700
Add opt-in concurrent pod creation to KubernetesExecutor (#68480)
---
providers/cncf/kubernetes/provider.yaml | 25 +
.../kubernetes/executors/kubernetes_executor.py | 272 ++++++++---
.../executors/kubernetes_executor_utils.py | 111 ++++-
.../providers/cncf/kubernetes/get_provider_info.py | 14 +
.../providers/cncf/kubernetes/hooks/kubernetes.py | 60 +--
.../providers/cncf/kubernetes/kube_client.py | 102 ++++
.../providers/cncf/kubernetes/kube_config.py | 7 +
.../cncf/kubernetes/kubernetes_helper_functions.py | 11 +-
.../executors/test_kubernetes_executor.py | 513 ++++++++++++++++++++-
.../tests/unit/cncf/kubernetes/test_kube_client.py | 39 ++
.../observability/metrics/metrics_template.yaml | 13 +
11 files changed, 1032 insertions(+), 135 deletions(-)
diff --git a/providers/cncf/kubernetes/provider.yaml b/providers/cncf/kubernetes/provider.yaml
index b9905d301ec..a6ba5b42ede 100644
--- a/providers/cncf/kubernetes/provider.yaml
+++ b/providers/cncf/kubernetes/provider.yaml
@@ -364,6 +364,31 @@ config:
type: string
example: ~
default: "1"
+ async_pod_creation:
+ description: |
+ Create worker pods concurrently within each scheduler loop instead of
+ sequentially. When enabled, the ``worker_pods_creation_batch_size`` pods
+ dequeued per loop are submitted to the Kubernetes API concurrently (bounded
+ by ``pod_creation_max_concurrency``) using the asynchronous Kubernetes client.
+ This reduces the time the scheduler loop spends blocked on pod creation when
+ per-call latency is high (network round-trip, admission webhooks). Pod
+ templates are still built synchronously; only the create API calls are
+ parallelized.
+ version_added: 10.20.0
+ type: boolean
+ example: ~
+ default: "False"
+ pod_creation_max_concurrency:
+ description: |
+ Maximum number of concurrent pod-creation API calls when ``async_pod_creation``
+ is enabled. Bounds the burst of simultaneous requests to the Kubernetes API
+ server (and admission webhooks) to avoid tripping API priority-and-fairness or
+ rate limits. Has no effect when ``async_pod_creation`` is False. When set to 0
+ the limit falls back to ``worker_pods_creation_batch_size``.
+ version_added: 10.20.0
+ type: integer
+ example: "16"
+ default: "0"
multi_namespace_mode:
description: |
Allows users to launch pods in multiple namespaces.
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 8aaa430b921..b75815df2fd 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
@@ -34,6 +34,7 @@ from collections.abc import Iterable
from contextlib import suppress
from dataclasses import dataclass
from datetime import datetime, timedelta
+from http import HTTPStatus
from itertools import chain
from queue import Empty, Queue
from typing import TYPE_CHECKING, Any
@@ -53,7 +54,10 @@ from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types impor
KubernetesResults,
)
from airflow.providers.cncf.kubernetes.kube_config import KubeConfig
-from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import annotations_to_key
+from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
+ TRANSIENT_CONNECTION_ERRORS,
+ 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.common.compat.sdk import Stats, conf
@@ -457,92 +461,210 @@ class KubernetesExecutor(BaseExecutor):
last_resource_version[ns] or resource_instance.resource_version[ns]
)
- from kubernetes.client.rest import ApiException
-
if self.create_pods_after and self.create_pods_after > datetime.now():
self.log.warning("Skipping pod creation due to kubernetes rate limit")
return
self.create_pods_after = None
+ if self.kube_config.async_pod_creation:
+ self._create_pods_concurrently()
+ else:
+ self._create_pods_sequentially()
+
+ def _create_pods_sequentially(self) -> None:
+ """Dequeue a batch and create worker pods one at a time (default behavior)."""
+ from kubernetes.client.rest import ApiException
+
+ if TYPE_CHECKING:
+ assert self.kube_scheduler
+ assert self.task_queue
+ created: int = 0
+ start: float = time.monotonic()
with contextlib.suppress(Empty):
for _ in range(self.kube_config.worker_pods_creation_batch_size):
- task = self.task_queue.get_nowait()
-
+ task: KubernetesJob = self.task_queue.get_nowait()
+ created += 1
try:
- key = task.key
self.kube_scheduler.run_next(task)
- self.task_publish_retries.pop(key, None)
- except PodReconciliationError as e:
- self.log.exception(
- "Pod reconciliation failed, likely due to kubernetes library upgrade. "
- "Try clearing the task to re-run.",
- )
- self.fail(task[0], e)
- except ApiException as e:
- try:
- if e.body:
- body = json.loads(e.body)
- else:
- # If no body content, use reason as the message
- body = {"message": e.reason}
- except (json.JSONDecodeError, ValueError, TypeError):
- # If the body is a string (e.g., in a 429 error), it can't be parsed as JSON.
- # Use the body directly as the message instead.
- body = {"message": e.body}
-
- headers = e.headers or {}
- retries = self.task_publish_retries[key]
- # In case of exceeded quota or conflict errors, requeue the task as per the task_publish_max_retries
- # In case of a rate limit, wait and do not create new pods for "Retry-After" seconds
- can_retry_publish = (
- self.task_publish_max_retries == -1 or retries < self.task_publish_max_retries
- )
- message = body.get("message", "")
- if (
- (str(e.status) == "403" and "exceeded quota" in message)
- or (str(e.status) == "409" and "object has been modified" in message)
- or (str(e.status) == "410" and "too old resource version" in message)
- or str(e.status) == "500"
- or str(e.status) == "429"
- ) and can_retry_publish:
- self.log.warning(
- "[Try %s of %s] Kube ApiException for Task: (%s). Reason: %r. Message: %s",
- self.task_publish_retries[key] + 1,
- self.task_publish_max_retries,
- key,
- e.reason,
- message,
- )
-
- self.task_queue.put(task)
- self.task_publish_retries[key] = retries + 1
-
- if str(e.status) == "429":
- self.create_pods_after = datetime.now() + timedelta(
- seconds=int(headers.get("Retry-After", "0"))
- )
- self.log.warning(
- "Got rate limit from k8s api, skipping pod creation until %s",
- self.create_pods_after,
- )
- # stop pod creation to stop api requests
- break
- else:
- self.log.error("Pod creation failed with reason %r. Failing task", e.reason)
- key = task.key
- self.fail(key, e)
- self.task_publish_retries.pop(key, None)
- except PodMutationHookException as e:
- key = task.key
- self.log.error(
- "Pod Mutation Hook failed for the task %s. Failing task. Details: %s",
- key,
- e.__cause__,
- )
- self.fail(key, e)
+ self.task_publish_retries.pop(task.key, None)
+ except (
+ PodReconciliationError,
+ ApiException,
+ PodMutationHookException,
+ *TRANSIENT_CONNECTION_ERRORS,
+ ) as e:
+ if self._handle_pod_publish_error(task, e):
+ # Rate limited: stop creating further pods this loop.
+ break
finally:
self.task_queue.task_done()
+ if created:
+ self._record_pod_creation_batch(created, time.monotonic() - start)
+
+ def _create_pods_concurrently(self) -> None:
+ """
+ Dequeue a batch and create worker pods concurrently via the async client.
+
+ The whole batch is in flight before any response returns, so (unlike the sequential
+ path) a mid-batch 429 cannot stop the burst — it is bounded only by
+ ``pod_creation_max_concurrency``, and a 429 still suppresses the next loop via
+ ``create_pods_after``.
+ """
+ if TYPE_CHECKING:
+ assert self.kube_scheduler
+ assert self.task_queue
+ 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())
+ if not jobs:
+ return
+ start: float = time.monotonic()
+ try:
+ results: list[tuple[KubernetesJob, Exception | None]] = self.kube_scheduler.run_next_batch(jobs)
+ except Exception:
+ # Catastrophic failure (e.g. async client setup): keep queue accounting correct
+ # by marking every dequeued task done before surfacing the error.
+ for _ in jobs:
+ self.task_queue.task_done()
+ raise
+ self._record_pod_creation_batch(len(jobs), time.monotonic() - start)
+ for task, error in results:
+ try:
+ if error is None:
+ self.task_publish_retries.pop(task.key, None)
+ else:
+ self._handle_pod_publish_error(task, error)
+ finally:
+ self.task_queue.task_done()
+
+ def _record_pod_creation_batch(self, count: int, duration_seconds: float) -> None:
+ """Emit per-loop pod-creation batch metrics (duration, size); both paths use the same names."""
+ Stats.timing("kubernetes_executor.pod_creation_batch_duration", timedelta(seconds=duration_seconds))
+ Stats.gauge("kubernetes_executor.pod_creation_batch_size", count)
+
+ def _handle_pod_publish_error(self, task: KubernetesJob, e: Exception) -> bool:
+ """
+ Handle a build/create failure for a worker pod; shared by both creation paths.
+
+ The sync and async clients' ApiExceptions expose the same fields, so they are handled
+ uniformly. Returns True if pod creation should stop for the rest of this scheduler loop
+ (rate limit), else False.
+ """
+ from kubernetes.client.rest import ApiException
+ from kubernetes_asyncio.client.exceptions import ApiException as AsyncApiException
+
+ if TYPE_CHECKING:
+ assert self.task_queue
+ key: TaskInstanceKey = task.key
+ if isinstance(e, PodReconciliationError):
+ self.log.exception(
+ "Pod reconciliation failed, likely due to kubernetes library upgrade. "
+ "Try clearing the task to re-run.",
+ )
+ self.fail(key, e)
+ return False
+ if isinstance(e, PodMutationHookException):
+ self.log.error(
+ "Pod Mutation Hook failed for the task %s. Failing task. Details: %s",
+ key,
+ e.__cause__,
+ )
+ self.fail(key, e)
+ return False
+ if isinstance(e, (ApiException, AsyncApiException)):
+ body: dict[str, Any]
+ try:
+ if e.body:
+ body = json.loads(e.body)
+ else:
+ # If no body content, use reason as the message
+ body = {"message": e.reason}
+ except (json.JSONDecodeError, ValueError, TypeError):
+ # If the body is a string (e.g., in a 429 error), it can't be parsed as JSON.
+ # Use the body directly as the message instead.
+ body = {"message": e.body}
+
+ headers = e.headers or {}
+ retries: int = self.task_publish_retries[key]
+ # Requeue transient failures (exceeded-quota / stale-version conflicts and the api
+ # server's 429 / 5xx) up to task_publish_max_retries; anything else fails immediately.
+ can_retry_publish = self.task_publish_max_retries == -1 or retries < self.task_publish_max_retries
+ message: str = body.get("message", "")
+ # Retry-After (seconds): the apiserver sets it on 429 (APF throttling) and on 503 while
+ # shutting down. Ignore a malformed value rather than crashing the scheduler loop.
+ retry_after = headers.get("Retry-After")
+ try:
+ retry_after_seconds: int | None = int(retry_after) if retry_after is not None else None
+ except (TypeError, ValueError):
+ retry_after_seconds = None
+ transient_5xx = (
+ HTTPStatus.INTERNAL_SERVER_ERROR,
+ HTTPStatus.BAD_GATEWAY,
+ HTTPStatus.SERVICE_UNAVAILABLE,
+ HTTPStatus.GATEWAY_TIMEOUT,
+ )
+ if (
+ (e.status == HTTPStatus.FORBIDDEN and "exceeded quota" in message)
+ or (e.status == HTTPStatus.CONFLICT and "object has been modified" in message)
+ or (e.status == HTTPStatus.GONE and "too old resource version" in message)
+ or e.status == HTTPStatus.TOO_MANY_REQUESTS
+ or e.status in transient_5xx
+ ) and can_retry_publish:
+ self.log.warning(
+ "[Try %s of %s] Kube ApiException for Task: (%s). Reason: %r. Message: %s",
+ self.task_publish_retries[key] + 1,
+ self.task_publish_max_retries,
+ key,
+ e.reason,
+ message,
+ )
+
+ self.task_queue.put(task)
+ self.task_publish_retries[key] = retries + 1
+
+ # Pause the loop until Retry-After when the server told us to back off (429, or 503
+ # on apiserver shutdown); other transient 5xx retry on the next loop.
+ if e.status == HTTPStatus.TOO_MANY_REQUESTS or retry_after_seconds is not None:
+ self.create_pods_after = datetime.now() + timedelta(seconds=retry_after_seconds or 0)
+ self.log.warning(
+ "Backing off pod creation until %s after status %s from k8s api",
+ self.create_pods_after,
+ e.status,
+ )
+ # stop pod creation to stop api requests
+ return True
+ else:
+ self.log.error("Pod creation failed with reason %r. Failing task", e.reason)
+ self.fail(key, e)
+ self.task_publish_retries.pop(key, None)
+ return False
+ if isinstance(e, TRANSIENT_CONNECTION_ERRORS):
+ # Connection-level failure talking to the api server (reset / DNS blip / read timeout) —
+ # the category generic_api_retry treats as transient. Re-queue rather than fail; the
+ # create may not have reached the server, so a later loop retries.
+ retries = self.task_publish_retries[key]
+ if self.task_publish_max_retries == -1 or retries < self.task_publish_max_retries:
+ self.log.warning(
+ "[Try %s of %s] Transient connection error creating pod for Task %s: %r",
+ retries + 1,
+ self.task_publish_max_retries,
+ key,
+ e,
+ )
+ self.task_queue.put(task)
+ self.task_publish_retries[key] = retries + 1
+ else:
+ self.log.error(
+ "Connection error creating pod, retries exhausted. Failing task %s: %r", key, e
+ )
+ self.fail(key, e)
+ self.task_publish_retries.pop(key, None)
+ return False
+ # Unknown exception type: fail the task rather than silently dropping it.
+ self.fail(key, e)
+ return False
@provide_session
def _change_state(
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 2c4bd306e52..0bf1ed8b23e 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
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+import asyncio
import contextlib
import json
import multiprocessing
@@ -25,6 +26,7 @@ from typing import TYPE_CHECKING, Any, Literal, cast
from kubernetes import client, watch
from kubernetes.client.rest import ApiException
+from kubernetes_asyncio import client as async_client
from urllib3.exceptions import ReadTimeoutError
from airflow.providers.cncf.kubernetes.backcompat import get_logical_date_key
@@ -38,7 +40,7 @@ from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types impor
KubernetesResults,
KubernetesWatch,
)
-from airflow.providers.cncf.kubernetes.kube_client import get_kube_client
+from airflow.providers.cncf.kubernetes.kube_client import get_async_kube_client, get_kube_client
from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
annotations_for_logging_task_metadata,
annotations_to_key,
@@ -51,6 +53,8 @@ from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.state import TaskInstanceState
if TYPE_CHECKING:
+ from collections.abc import Iterator
+
from kubernetes.client import Configuration, models as k8s
@@ -489,6 +493,12 @@ class AirflowKubernetesScheduler(LoggingMixin):
self.scheduler_job_id = scheduler_job_id
self.kube_watchers = self._make_kube_watchers()
self.team_name = team_name
+ # Async pod-creation state; populated lazily, only used when async_pod_creation is enabled.
+ self._async_loop: asyncio.AbstractEventLoop | None = None
+ self._async_pod_client: async_client.CoreV1Api | None = None
+ self.pod_creation_max_concurrency = (
+ self.kube_config.pod_creation_max_concurrency or self.kube_config.worker_pods_creation_batch_size
+ )
def run_pod_async(self, pod: k8s.V1Pod, **kwargs):
"""Run POD asynchronously."""
@@ -568,6 +578,18 @@ class AirflowKubernetesScheduler(LoggingMixin):
def run_next(self, next_job: KubernetesJob) -> None:
"""Receives the next job to run, builds the pod, and creates it."""
+ pod = self._build_pod_request(next_job)
+ # the watcher will monitor pods, so we do not block.
+ self.run_pod_async(pod, **self.kube_config.kube_client_request_args)
+ self.log.debug("Kubernetes Job created!")
+
+ def _build_pod_request(self, next_job: KubernetesJob) -> k8s.V1Pod:
+ """
+ Build the worker pod request object for a job.
+
+ Performs no API calls. May raise ``PodMutationHookException`` or
+ ``PodReconciliationError`` from the pod-mutation hook / reconciliation.
+ """
key = next_job.key
command = next_job.command
kube_executor_config = next_job.kube_executor_config
@@ -621,10 +643,90 @@ class AirflowKubernetesScheduler(LoggingMixin):
)
self.log.debug("Kubernetes running for command %s", command)
self.log.debug("Kubernetes launching image %s", pod.spec.containers[0].image)
+ return pod
- # the watcher will monitor pods, so we do not block.
- self.run_pod_async(pod, **self.kube_config.kube_client_request_args)
- self.log.debug("Kubernetes Job created!")
+ def run_next_batch(self, next_jobs: list[KubernetesJob]) -> list[tuple[KubernetesJob, Exception | None]]:
+ """
+ Build pod requests synchronously, then create them concurrently via the async client.
+
+ Bounded by ``pod_creation_max_concurrency``. Returns one ``(job, exception)`` per job —
+ build and create failures are returned, not raised, so the caller handles them like the
+ sequential path.
+ """
+ built: list[tuple[KubernetesJob, k8s.V1Pod | None, Exception | None]] = []
+ for job in next_jobs:
+ try:
+ built.append((job, self._build_pod_request(job), None))
+ except Exception as e:
+ built.append((job, None, e))
+
+ to_create: list[tuple[KubernetesJob, k8s.V1Pod]] = [
+ (job, pod) for job, pod, build_err in built if build_err is None and pod is not None
+ ]
+ create_errors: list[Exception | None] = self._run_pods_async(to_create) if to_create else []
+
+ # create_errors aligns 1:1 with to_create (gather preserves order); walk it as we
+ # re-emit one (job, error) per built entry, pairing build failures with their own error.
+ create_iter: Iterator[Exception | None] = iter(create_errors)
+ results: list[tuple[KubernetesJob, Exception | None]] = []
+ for job, _, build_err in built:
+ results.append((job, build_err if build_err is not None else next(create_iter)))
+ return results
+
+ def _run_pods_async(self, jobs_and_pods: list[tuple[KubernetesJob, k8s.V1Pod]]) -> list[Exception | None]:
+ """Create the given pods concurrently on a dedicated event loop; one error (or None) per pod, in order."""
+ if self._async_loop is None:
+ self._async_loop = asyncio.new_event_loop()
+ return self._async_loop.run_until_complete(self._create_pods_async(jobs_and_pods))
+
+ async def _create_pods_async(
+ self, jobs_and_pods: list[tuple[KubernetesJob, k8s.V1Pod]]
+ ) -> list[Exception | None]:
+ """Issue create_namespaced_pod calls concurrently, bounded by a semaphore; one result per pod, in order."""
+ if self._async_pod_client is None:
+ self._async_pod_client = await get_async_kube_client()
+ api = self._async_pod_client
+ semaphore = asyncio.Semaphore(self.pod_creation_max_concurrency)
+ request_kwargs: dict[str, Any] = self.kube_config.kube_client_request_args or {}
+
+ async def _create(pod: k8s.V1Pod) -> None:
+ # Sanitize with the sync client (identical to run_pod_async) to guarantee the
+ # request body matches the sequential path exactly.
+ sanitized_pod = self.kube_client.api_client.sanitize_for_serialization(pod)
+ async with semaphore:
+ try:
+ with Stats.timer("kubernetes_executor.pod_creation"):
+ await api.create_namespaced_pod(
+ body=sanitized_pod, namespace=pod.metadata.namespace, **request_kwargs
+ )
+ Stats.incr("kubernetes_executor.pod_creation_status", tags={"status": "200"})
+ except async_client.exceptions.ApiException as e:
+ Stats.incr("kubernetes_executor.pod_creation_status", tags={"status": str(e.status)})
+ raise
+ except Exception:
+ Stats.incr("kubernetes_executor.pod_creation_status", tags={"status": "error"})
+ raise
+
+ outcomes: list[BaseException | None] = await asyncio.gather(
+ *(_create(pod) for _, pod in jobs_and_pods), return_exceptions=True
+ )
+ return [outcome if isinstance(outcome, Exception) else None for outcome in outcomes]
+
+ def _close_async_pod_client(self) -> None:
+ """Close the async pod client and its event loop, if they were created."""
+ if self._async_loop is None:
+ return
+ if self._async_pod_client is not None:
+ try:
+ # CoreV1Api sets self.api_client in __init__, but kubernetes_asyncio's
+ # swagger-codegen stubs don't expose it at the class level.
+ api_client = self._async_pod_client.api_client # type: ignore[attr-defined]
+ self._async_loop.run_until_complete(api_client.close())
+ except Exception:
+ self.log.warning("Error while closing async pod client", exc_info=True)
+ self._async_loop.close()
+ self._async_loop = None
+ self._async_pod_client = None
def delete_pod(self, pod_name: str, namespace: str) -> None:
"""Delete Pod from a namespace; does not raise if it does not exist."""
@@ -756,6 +858,7 @@ class AirflowKubernetesScheduler(LoggingMixin):
def terminate(self) -> None:
"""Terminates the watcher."""
+ self._close_async_pod_client()
self.log.debug("Terminating kube_watchers...")
for kube_watcher in self.kube_watchers.values():
kube_watcher.terminate()
diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py
index df11e510921..8a2daec956f 100644
--- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py
+++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py
@@ -222,6 +222,20 @@ def get_provider_info():
"example": None,
"default": "1",
},
+ "async_pod_creation": {
+ "description": "Create worker pods concurrently within each scheduler loop instead of\nsequentially. When enabled, the ``worker_pods_creation_batch_size`` pods\ndequeued per loop are submitted to the Kubernetes API concurrently (bounded\nby ``pod_creation_max_concurrency``) using the asynchronous Kubernetes client.\nThis reduces the time the scheduler loop spends blocked on pod creation when\nper-call latency is high (network round-trip, admission webhooks). Pod\n [...]
+ "version_added": "10.20.0",
+ "type": "boolean",
+ "example": None,
+ "default": "False",
+ },
+ "pod_creation_max_concurrency": {
+ "description": "Maximum number of concurrent pod-creation API calls when ``async_pod_creation``\nis enabled. Bounds the burst of simultaneous requests to the Kubernetes API\nserver (and admission webhooks) to avoid tripping API priority-and-fairness or\nrate limits. Has no effect when ``async_pod_creation`` is False. When set to 0\nthe limit falls back to ``worker_pods_creation_batch_size``.\n",
+ "version_added": "10.20.0",
+ "type": "integer",
+ "example": "16",
+ "default": "0",
+ },
"multi_namespace_mode": {
"description": "Allows users to launch pods in multiple namespaces.\nWill require creating a cluster-role for the scheduler,\nor use multi_namespace_mode_namespace_list configuration.\n",
"version_added": None,
diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py
index 7c8140a6f07..8ebf00c69e4 100644
--- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py
+++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py
@@ -35,12 +35,13 @@ from urllib3.exceptions import HTTPError
from airflow.models import Connection
from airflow.providers.cncf.kubernetes.exceptions import KubernetesApiError, KubernetesApiPermissionError
-from airflow.providers.cncf.kubernetes.kube_client import _disable_verify_ssl, _enable_tcp_keepalive
-from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
- API_TIMEOUT,
- API_TIMEOUT_OFFSET_SERVER_SIDE,
- generic_api_retry,
+from airflow.providers.cncf.kubernetes.kube_client import (
+ _disable_verify_ssl,
+ _enable_tcp_keepalive,
+ _TimeoutAsyncK8sApiClient,
+ _TimeoutK8sApiClient,
)
+from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import generic_api_retry
from airflow.providers.cncf.kubernetes.utils.container import (
container_is_completed,
container_is_running,
@@ -74,55 +75,6 @@ def _load_body_to_dict(body: str) -> dict:
return body_dict
-def _get_request_timeout(timeout_seconds: int | None) -> float:
- """Get the client-side request timeout."""
- if timeout_seconds is not None and timeout_seconds > API_TIMEOUT - API_TIMEOUT_OFFSET_SERVER_SIDE:
- return timeout_seconds + API_TIMEOUT_OFFSET_SERVER_SIDE
- return API_TIMEOUT
-
-
-class _TimeoutK8sApiClient(client.ApiClient):
- """
- Wrapper around kubernetes sync ApiClient to set default timeout.
-
- When *disable_verify_ssl* is True the TLS certificate check is turned off
- on the *client_configuration* that is passed (or on a fresh default copy)
- so that callers do not need to repeat this logic at every call-site.
- """
-
- def __init__(
- self,
- configuration: client.Configuration | None = None,
- *,
- disable_verify_ssl: bool = False,
- ) -> None:
- if disable_verify_ssl:
- if configuration is None:
- configuration = client.Configuration.get_default_copy()
- configuration.verify_ssl = False
- super().__init__(configuration=configuration)
-
- def call_api(self, *args, **kwargs):
- timeout_seconds = kwargs.get("timeout_seconds") # get server-side timeout
- kwargs.setdefault("_request_timeout", _get_request_timeout(timeout_seconds)) # client-side timeout
- return super().call_api(*args, **kwargs)
-
-
-class _TimeoutAsyncK8sApiClient(async_client.ApiClient):
- """Wrapper around kubernetes async ApiClient to set default timeout."""
-
- def __init__(
- self,
- configuration: async_client.Configuration | None = None,
- ) -> None:
- super().__init__(configuration=configuration)
-
- async def call_api(self, *args, **kwargs):
- timeout_seconds = kwargs.get("timeout_seconds") # server-side timeout
- kwargs.setdefault("_request_timeout", _get_request_timeout(timeout_seconds)) # client-side timeout
- return await super().call_api(*args, **kwargs)
-
-
class PodOperatorHookProtocol(Protocol):
"""
Protocol to define methods relied upon by KubernetesPodOperator.
diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_client.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_client.py
index a4c5f405577..79e2b802e01 100644
--- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_client.py
+++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_client.py
@@ -19,6 +19,7 @@
from __future__ import annotations
import logging
+from typing import Any
import urllib3.util
@@ -30,6 +31,12 @@ try:
from kubernetes import client, config
from kubernetes.client import Configuration
from kubernetes.client.rest import ApiException
+ from kubernetes_asyncio import client as async_client, config as async_config
+
+ from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
+ API_TIMEOUT,
+ API_TIMEOUT_OFFSET_SERVER_SIDE,
+ )
has_kubernetes = True
@@ -43,6 +50,56 @@ try:
configuration.verify_ssl = False
Configuration.set_default(configuration)
+ def _get_request_timeout(timeout_seconds: int | None) -> float:
+ """Get the client-side request timeout."""
+ if timeout_seconds is not None and timeout_seconds > API_TIMEOUT - API_TIMEOUT_OFFSET_SERVER_SIDE:
+ return timeout_seconds + API_TIMEOUT_OFFSET_SERVER_SIDE
+ return API_TIMEOUT
+
+ class _TimeoutK8sApiClient(client.ApiClient):
+ """
+ Wrapper around kubernetes sync ApiClient to set default timeout.
+
+ When *disable_verify_ssl* is True the TLS certificate check is turned off
+ on the *client_configuration* that is passed (or on a fresh default copy)
+ so that callers do not need to repeat this logic at every call-site.
+ """
+
+ def __init__(
+ self,
+ configuration: client.Configuration | None = None,
+ *,
+ disable_verify_ssl: bool = False,
+ ) -> None:
+ if disable_verify_ssl:
+ if configuration is None:
+ configuration = client.Configuration.get_default_copy()
+ configuration.verify_ssl = False
+ super().__init__(configuration=configuration)
+
+ def call_api(self, *args: Any, **kwargs: Any) -> Any:
+ timeout_seconds = kwargs.get("timeout_seconds") # get server-side timeout
+ kwargs.setdefault(
+ "_request_timeout", _get_request_timeout(timeout_seconds)
+ ) # client-side timeout
+ return super().call_api(*args, **kwargs)
+
+ class _TimeoutAsyncK8sApiClient(async_client.ApiClient):
+ """Wrapper around kubernetes async ApiClient to set default timeout."""
+
+ def __init__(
+ self,
+ configuration: async_client.Configuration | None = None,
+ ) -> None:
+ super().__init__(configuration=configuration)
+
+ async def call_api(self, *args: Any, **kwargs: Any) -> Any:
+ timeout_seconds = kwargs.get("timeout_seconds") # server-side timeout
+ kwargs.setdefault(
+ "_request_timeout", _get_request_timeout(timeout_seconds)
+ ) # client-side timeout
+ return await super().call_api(*args, **kwargs)
+
except ImportError as e:
# We need an exception class to be able to use it in ``except`` elsewhere
# in the code base
@@ -154,3 +211,48 @@ def get_kube_client(
api_client = client.ApiClient(configuration=configuration)
return client.CoreV1Api(api_client)
+
+
+async def get_async_kube_client(
+ in_cluster: bool | None = None,
+ cluster_context: str | None = None,
+ config_file: str | None = None,
+) -> async_client.CoreV1Api:
+ """
+ Retrieve an asynchronous Kubernetes client.
+
+ Mirrors :func:`get_kube_client` but builds a ``kubernetes_asyncio`` client so that
+ pod-creation API calls can be issued concurrently by the KubernetesExecutor. Reads the
+ same configuration keys (``in_cluster``, ``cluster_context``, ``config_file``,
+ ``verify_ssl``, ``ssl_ca_cert``).
+
+ :param in_cluster: whether we are in cluster
+ :param cluster_context: context of the cluster
+ :param config_file: configuration file
+ :return: asynchronous kubernetes client
+ """
+ if not has_kubernetes:
+ raise _import_err
+ if in_cluster is None:
+ in_cluster = conf.getboolean("kubernetes_executor", "in_cluster")
+
+ configuration = async_client.Configuration()
+ if not conf.getboolean("kubernetes_executor", "verify_ssl"):
+ configuration.verify_ssl = False
+
+ if in_cluster:
+ async_config.load_incluster_config(client_configuration=configuration)
+ else:
+ if cluster_context is None:
+ cluster_context = conf.get("kubernetes_executor", "cluster_context", fallback=None)
+ if config_file is None:
+ config_file = conf.get("kubernetes_executor", "config_file", fallback=None)
+ await async_config.load_kube_config(
+ config_file=config_file, context=cluster_context, client_configuration=configuration
+ )
+
+ ssl_ca_cert = conf.get("kubernetes_executor", "ssl_ca_cert")
+ if ssl_ca_cert:
+ configuration.ssl_ca_cert = ssl_ca_cert
+
+ return async_client.CoreV1Api(_TimeoutAsyncK8sApiClient(configuration))
diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_config.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_config.py
index cc9d7fc08fe..f9462bdaa99 100644
--- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_config.py
+++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_config.py
@@ -67,6 +67,13 @@ class KubeConfig:
self.worker_pods_creation_batch_size = self._conf.getint(
self.kubernetes_section, "worker_pods_creation_batch_size"
)
+ self.async_pod_creation = self._conf.getboolean(
+ self.kubernetes_section, "async_pod_creation", fallback=False
+ )
+ # 0 means "fall back to worker_pods_creation_batch_size" (resolved in the scheduler).
+ self.pod_creation_max_concurrency = self._conf.getint(
+ self.kubernetes_section, "pod_creation_max_concurrency", fallback=0
+ )
self.worker_container_repository = self._conf.get(
self.kubernetes_section, "worker_container_repository"
)
diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py
index 5a41da5478f..9b32fb5967f 100644
--- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py
+++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py
@@ -24,6 +24,7 @@ from typing import TYPE_CHECKING
import pendulum
import tenacity
+from aiohttp import ClientConnectionError
from kubernetes.client.rest import ApiException as SyncApiException
from kubernetes_asyncio.client.exceptions import ApiException as AsyncApiException
from slugify import slugify
@@ -60,13 +61,21 @@ API_RETRY_WAIT_MAX = conf.getfloat("workers", "api_retry_wait_max", fallback=15)
_default_wait = tenacity.wait_exponential(min=API_RETRY_WAIT_MIN, max=API_RETRY_WAIT_MAX)
TRANSIENT_STATUS_CODES = {409, 429, 500, 502, 503, 504}
+# Connection-level failures (socket reset, DNS blip, read timeout) worth retrying — the api server
+# never saw the request, or its reply was lost. Covers both kube clients: urllib3 (sync) and aiohttp
+# (async). Shared so this decorator and the KubernetesExecutor agree on what counts as transient.
+TRANSIENT_CONNECTION_ERRORS: tuple[type[BaseException], ...] = (
+ HTTPError,
+ ClientConnectionError,
+ KubernetesApiException,
+)
def _should_retry_api(exc: BaseException) -> bool:
"""Retry on selected ApiException status codes, plus plain HTTP/timeout errors."""
if isinstance(exc, (SyncApiException, AsyncApiException)):
return exc.status in TRANSIENT_STATUS_CODES
- return isinstance(exc, (HTTPError, KubernetesApiException))
+ return isinstance(exc, TRANSIENT_CONNECTION_ERRORS)
class WaitRetryAfterOrExponential(tenacity.wait.wait_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 ee541bd83fd..5f1c03e7d45 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
@@ -21,14 +21,17 @@ import re
import string
import time
from datetime import datetime, timedelta
+from queue import Queue
from unittest import mock
import pytest
import yaml
+from aiohttp import ClientConnectionError
from kubernetes.client import models as k8s
from kubernetes.client.rest import ApiException
from sqlalchemy import inspect
-from urllib3 import HTTPResponse
+from urllib3 import HTTPConnectionPool, HTTPResponse
+from urllib3.exceptions import MaxRetryError, ProtocolError
from airflow.jobs.job import Job
from airflow.models.taskinstancekey import TaskInstanceKey
@@ -877,6 +880,46 @@ class TestKubernetesExecutor:
None,
id="500 Internal Server Error (webhook failure) (retry failed)",
),
+ pytest.param(
+ HTTPResponse(body='{"message": "Bad Gateway"}', status=502),
+ 1,
+ True,
+ State.SUCCESS,
+ None,
+ id="502 Bad Gateway (transient, requeued, retry next loop)",
+ ),
+ pytest.param(
+ HTTPResponse(
+ body='{"message": "apiserver is shutting down"}',
+ status=503,
+ headers={"Retry-After": "1"},
+ ),
+ 1,
+ True,
+ State.SUCCESS,
+ 1,
+ id="503 Service Unavailable (Retry-After honored, requeued after retry delay)",
+ ),
+ pytest.param(
+ HTTPResponse(body='{"message": "Gateway Timeout"}', status=504),
+ 1,
+ True,
+ State.SUCCESS,
+ None,
+ id="504 Gateway Timeout (transient, requeued, retry next loop)",
+ ),
+ pytest.param(
+ HTTPResponse(
+ body='{"message": "apiserver is shutting down"}',
+ status=503,
+ headers={"Retry-After": "1"},
+ ),
+ 1,
+ True,
+ State.FAILED,
+ 1,
+ id="503 Service Unavailable (retries exhausted, failed)",
+ ),
],
)
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
@@ -1126,6 +1169,474 @@ class TestKubernetesExecutor:
finally:
kubernetes_executor.end()
+ @pytest.mark.skipif(
+ AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
+ )
+ @pytest.mark.parametrize(
+ ("exc", "task_publish_max_retries", "should_requeue"),
+ [
+ pytest.param(ProtocolError("Connection aborted."), 1, True, id="connection reset (requeued)"),
+ pytest.param(
+ MaxRetryError(
+ HTTPConnectionPool("localhost"), "/api/v1/namespaces/default/pods", Exception("refused")
+ ),
+ 1,
+ True,
+ id="client connect retries exhausted (requeued)",
+ ),
+ pytest.param(
+ ProtocolError("Connection aborted."), 0, False, id="connection reset, retries off (failed)"
+ ),
+ ],
+ )
+ @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ def test_run_next_connection_error_requeue(
+ self,
+ mock_get_kube_client,
+ mock_kubernetes_job_watcher,
+ exc,
+ task_publish_max_retries,
+ should_requeue,
+ data_file,
+ ):
+ """A transient connection error on the sync client re-queues the task instead of failing it."""
+ template_file = data_file("pods/generator_base_with_secrets.yaml").as_posix()
+
+ mock_kube_client = mock.patch("kubernetes.client.CoreV1Api", autospec=True)
+ mock_kube_client.create_namespaced_pod = mock.MagicMock(side_effect=exc)
+ mock_get_kube_client.return_value = mock_kube_client
+ mock_api_client = mock.MagicMock()
+ mock_api_client.sanitize_for_serialization.return_value = {}
+ mock_kube_client.api_client = mock_api_client
+
+ config = {("kubernetes_executor", "pod_template_file"): template_file}
+ with conf_vars(config):
+ kubernetes_executor = self.kubernetes_executor
+ kubernetes_executor.task_publish_max_retries = task_publish_max_retries
+ kubernetes_executor.start()
+ try:
+ task_instance_key = TaskInstanceKey("dag", "task", "run_id", 1)
+ kubernetes_executor.execute_async(
+ key=task_instance_key,
+ queue=None,
+ command=["airflow", "tasks", "run", "true", "some_parameter"],
+ )
+ kubernetes_executor.sync()
+
+ assert mock_kube_client.create_namespaced_pod.call_count == 1
+ if should_requeue:
+ assert not kubernetes_executor.task_queue.empty()
+ else:
+ assert kubernetes_executor.task_queue.empty()
+ assert kubernetes_executor.event_buffer[task_instance_key][0] == State.FAILED
+ finally:
+ kubernetes_executor.end()
+
+ @pytest.mark.skipif(
+ AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
+ )
+ @mock.patch(
+ "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+ new_callable=mock.AsyncMock,
+ )
+ @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ def test_async_pod_creation_creates_all_pods(
+ self, mock_get_kube_client, mock_kubernetes_job_watcher, mock_get_async_client, data_file
+ ):
+ """With async_pod_creation enabled, every dequeued pod is created via the async client."""
+ template_file = data_file("pods/generator_base_with_secrets.yaml").as_posix()
+ mock_kube_client = mock.MagicMock()
+ mock_kube_client.api_client.sanitize_for_serialization.return_value = {}
+ mock_get_kube_client.return_value = mock_kube_client
+
+ mock_async_api = mock.MagicMock()
+ mock_async_api.create_namespaced_pod = mock.AsyncMock()
+ mock_async_api.api_client.close = mock.AsyncMock()
+ mock_get_async_client.return_value = mock_async_api
+
+ config = {
+ ("kubernetes_executor", "pod_template_file"): template_file,
+ ("kubernetes_executor", "async_pod_creation"): "True",
+ ("kubernetes_executor", "pod_creation_max_concurrency"): "0",
+ ("kubernetes_executor", "worker_pods_creation_batch_size"): "16",
+ }
+ with conf_vars(config):
+ executor = KubernetesExecutor()
+ executor.job_id = 5
+ executor._last_completed_pod_adoption = time.monotonic()
+ executor.start()
+ try:
+ for i in range(5):
+ executor.execute_async(
+ key=TaskInstanceKey("dag", f"task{i}", "run_id", 1),
+ queue=None,
+ command=["airflow", "tasks", "run", "true", "x"],
+ )
+ executor.sync()
+ assert mock_async_api.create_namespaced_pod.await_count == 5
+ assert executor.task_queue.empty()
+ # max_concurrency=0 falls back to worker_pods_creation_batch_size
+ assert executor.kube_scheduler.pod_creation_max_concurrency == 16
+ finally:
+ executor.end()
+
+ @pytest.mark.skipif(
+ AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
+ )
+ @mock.patch(
+ "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+ new_callable=mock.AsyncMock,
+ )
+ @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ def test_async_pod_creation_respects_concurrency_limit(
+ self, mock_get_kube_client, mock_kubernetes_job_watcher, mock_get_async_client, data_file
+ ):
+ """The number of in-flight create calls never exceeds pod_creation_max_concurrency."""
+ import asyncio as _asyncio
+
+ template_file = data_file("pods/generator_base_with_secrets.yaml").as_posix()
+ mock_kube_client = mock.MagicMock()
+ mock_kube_client.api_client.sanitize_for_serialization.return_value = {}
+ mock_get_kube_client.return_value = mock_kube_client
+
+ inflight = {"current": 0, "peak": 0}
+
+ async def _tracked_create(*args, **kwargs):
+ inflight["current"] += 1
+ inflight["peak"] = max(inflight["peak"], inflight["current"])
+ await _asyncio.sleep(0.02)
+ inflight["current"] -= 1
+
+ mock_async_api = mock.MagicMock()
+ mock_async_api.create_namespaced_pod = mock.AsyncMock(side_effect=_tracked_create)
+ mock_async_api.api_client.close = mock.AsyncMock()
+ mock_get_async_client.return_value = mock_async_api
+
+ config = {
+ ("kubernetes_executor", "pod_template_file"): template_file,
+ ("kubernetes_executor", "async_pod_creation"): "True",
+ ("kubernetes_executor", "pod_creation_max_concurrency"): "3",
+ ("kubernetes_executor", "worker_pods_creation_batch_size"): "16",
+ }
+ with conf_vars(config):
+ executor = KubernetesExecutor()
+ executor.job_id = 5
+ executor._last_completed_pod_adoption = time.monotonic()
+ executor.start()
+ try:
+ for i in range(9):
+ executor.execute_async(
+ key=TaskInstanceKey("dag", f"task{i}", "run_id", 1),
+ queue=None,
+ command=["airflow", "tasks", "run", "true", "x"],
+ )
+ executor.sync()
+ assert mock_async_api.create_namespaced_pod.await_count == 9
+ assert inflight["peak"] <= 3
+ finally:
+ executor.end()
+
+ @pytest.mark.skipif(
+ AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
+ )
+ @mock.patch(
+ "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+ new_callable=mock.AsyncMock,
+ )
+ @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ def test_async_pod_creation_requeues_on_exceeded_quota(
+ self, mock_get_kube_client, mock_kubernetes_job_watcher, mock_get_async_client, data_file
+ ):
+ """An async-client quota error requeues the task via the shared pod-publish error handler."""
+ from kubernetes_asyncio.client.exceptions import ApiException as AsyncApiException
+
+ template_file = data_file("pods/generator_base_with_secrets.yaml").as_posix()
+ mock_kube_client = mock.MagicMock()
+ mock_kube_client.api_client.sanitize_for_serialization.return_value = {}
+ mock_get_kube_client.return_value = mock_kube_client
+
+ quota_exc = AsyncApiException(status=403, reason="Forbidden")
+ quota_exc.body = '{"message": "pods \\"x\\" is forbidden: exceeded quota: my-quota"}'
+ quota_exc.headers = {}
+
+ mock_async_api = mock.MagicMock()
+ mock_async_api.create_namespaced_pod = mock.AsyncMock(side_effect=quota_exc)
+ mock_async_api.api_client.close = mock.AsyncMock()
+ mock_get_async_client.return_value = mock_async_api
+
+ config = {
+ ("kubernetes_executor", "pod_template_file"): template_file,
+ ("kubernetes_executor", "async_pod_creation"): "True",
+ }
+ with conf_vars(config):
+ executor = KubernetesExecutor()
+ executor.job_id = 5
+ executor._last_completed_pod_adoption = time.monotonic()
+ executor.task_publish_max_retries = 1
+ executor.start()
+ try:
+ executor.execute_async(
+ key=TaskInstanceKey("dag", "task", "run_id", 1),
+ queue=None,
+ command=["airflow", "tasks", "run", "true", "x"],
+ )
+ executor.sync()
+ assert mock_async_api.create_namespaced_pod.await_count == 1
+ # Quota error is retryable -> task is requeued rather than failed.
+ assert not executor.task_queue.empty()
+ finally:
+ executor.end()
+
+ @pytest.mark.skipif(
+ AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
+ )
+ @mock.patch(
+ "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+ new_callable=mock.AsyncMock,
+ )
+ @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ def test_async_pod_creation_rate_limit_sets_create_pods_after(
+ self, mock_get_kube_client, mock_kubernetes_job_watcher, mock_get_async_client, data_file
+ ):
+ """A 429 from the async client sets create_pods_after so the next loop backs off."""
+ from kubernetes_asyncio.client.exceptions import ApiException as AsyncApiException
+
+ template_file = data_file("pods/generator_base_with_secrets.yaml").as_posix()
+ mock_kube_client = mock.MagicMock()
+ mock_kube_client.api_client.sanitize_for_serialization.return_value = {}
+ mock_get_kube_client.return_value = mock_kube_client
+
+ rate_exc = AsyncApiException(status=429, reason="Too Many Requests")
+ rate_exc.body = '{"message": "slow down"}'
+ rate_exc.headers = {"Retry-After": "1"}
+
+ mock_async_api = mock.MagicMock()
+ mock_async_api.create_namespaced_pod = mock.AsyncMock(side_effect=rate_exc)
+ mock_async_api.api_client.close = mock.AsyncMock()
+ mock_get_async_client.return_value = mock_async_api
+
+ config = {
+ ("kubernetes_executor", "pod_template_file"): template_file,
+ ("kubernetes_executor", "async_pod_creation"): "True",
+ }
+ with conf_vars(config):
+ executor = KubernetesExecutor()
+ executor.job_id = 5
+ executor._last_completed_pod_adoption = time.monotonic()
+ executor.task_publish_max_retries = 1
+ executor.start()
+ try:
+ executor.execute_async(
+ key=TaskInstanceKey("dag", "task", "run_id", 1),
+ queue=None,
+ command=["airflow", "tasks", "run", "true", "x"],
+ )
+ executor.sync()
+ assert executor.create_pods_after is not None
+ assert not executor.task_queue.empty()
+ finally:
+ executor.end()
+
+ @pytest.mark.skipif(
+ AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
+ )
+ @mock.patch(
+ "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+ new_callable=mock.AsyncMock,
+ )
+ @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ def test_async_pod_creation_requeues_and_backs_off_on_503(
+ self, mock_get_kube_client, mock_kubernetes_job_watcher, mock_get_async_client, data_file
+ ):
+ """A 503 + Retry-After from the async client (apiserver shutting down) requeues and backs off."""
+ from kubernetes_asyncio.client.exceptions import ApiException as AsyncApiException
+
+ template_file = data_file("pods/generator_base_with_secrets.yaml").as_posix()
+ mock_kube_client = mock.MagicMock()
+ mock_kube_client.api_client.sanitize_for_serialization.return_value = {}
+ mock_get_kube_client.return_value = mock_kube_client
+
+ unavailable_exc = AsyncApiException(status=503, reason="Service Unavailable")
+ unavailable_exc.body = '{"message": "apiserver is shutting down"}'
+ unavailable_exc.headers = {"Retry-After": "1"}
+
+ mock_async_api = mock.MagicMock()
+ mock_async_api.create_namespaced_pod = mock.AsyncMock(side_effect=unavailable_exc)
+ mock_async_api.api_client.close = mock.AsyncMock()
+ mock_get_async_client.return_value = mock_async_api
+
+ config = {
+ ("kubernetes_executor", "pod_template_file"): template_file,
+ ("kubernetes_executor", "async_pod_creation"): "True",
+ }
+ with conf_vars(config):
+ executor = KubernetesExecutor()
+ executor.job_id = 5
+ executor._last_completed_pod_adoption = time.monotonic()
+ executor.task_publish_max_retries = 1
+ executor.start()
+ try:
+ executor.execute_async(
+ key=TaskInstanceKey("dag", "task", "run_id", 1),
+ queue=None,
+ command=["airflow", "tasks", "run", "true", "x"],
+ )
+ executor.sync()
+ assert mock_async_api.create_namespaced_pod.await_count == 1
+ assert executor.create_pods_after is not None
+ assert not executor.task_queue.empty()
+ finally:
+ executor.end()
+
+ @pytest.mark.skipif(
+ AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
+ )
+ @mock.patch(
+ "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+ new_callable=mock.AsyncMock,
+ )
+ @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ def test_async_pod_creation_requeues_on_connection_error(
+ self, mock_get_kube_client, mock_kubernetes_job_watcher, mock_get_async_client, data_file
+ ):
+ """A connection error from the async (aiohttp) client requeues the task via the shared handler."""
+ template_file = data_file("pods/generator_base_with_secrets.yaml").as_posix()
+ mock_kube_client = mock.MagicMock()
+ mock_kube_client.api_client.sanitize_for_serialization.return_value = {}
+ mock_get_kube_client.return_value = mock_kube_client
+
+ mock_async_api = mock.MagicMock()
+ mock_async_api.create_namespaced_pod = mock.AsyncMock(
+ side_effect=ClientConnectionError("Cannot connect to host")
+ )
+ mock_async_api.api_client.close = mock.AsyncMock()
+ mock_get_async_client.return_value = mock_async_api
+
+ config = {
+ ("kubernetes_executor", "pod_template_file"): template_file,
+ ("kubernetes_executor", "async_pod_creation"): "True",
+ }
+ with conf_vars(config):
+ executor = KubernetesExecutor()
+ executor.job_id = 5
+ executor._last_completed_pod_adoption = time.monotonic()
+ executor.task_publish_max_retries = 1
+ executor.start()
+ try:
+ executor.execute_async(
+ key=TaskInstanceKey("dag", "task", "run_id", 1),
+ queue=None,
+ command=["airflow", "tasks", "run", "true", "x"],
+ )
+ executor.sync()
+ assert mock_async_api.create_namespaced_pod.await_count == 1
+ # Connection error is transient -> requeued (not failed).
+ assert not executor.task_queue.empty()
+ finally:
+ executor.end()
+
+ @pytest.mark.skipif(
+ AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
+ )
+ @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ def test_concurrent_path_fails_errored_and_completes_successful(
+ self, mock_get_kube_client, mock_kubernetes_job_watcher
+ ):
+ """A build/create failure for one job in a batch must not affect the others (isolation)."""
+ from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types import KubernetesJob
+
+ executor = self.kubernetes_executor
+ executor.kube_config.worker_pods_creation_batch_size = 16
+ executor.kube_scheduler = mock.MagicMock()
+ # start() is not called here, so the lazily-created task_queue (None until start) is
+ # supplied directly — this unit-tests _create_pods_concurrently's error isolation without
+ # spinning up the executor's Manager process.
+ executor.task_queue = Queue()
+ ok_key = TaskInstanceKey("dag", "ok", "run_id", 1)
+ bad_key = TaskInstanceKey("dag", "bad", "run_id", 1)
+ ok_job = KubernetesJob(ok_key, ["airflow", "tasks", "run"], {}, None)
+ bad_job = KubernetesJob(bad_key, ["airflow", "tasks", "run"], {}, None)
+ executor.task_queue.put(ok_job)
+ executor.task_queue.put(bad_job)
+ executor.kube_scheduler.run_next_batch.return_value = [
+ (ok_job, None),
+ (bad_job, PodReconciliationError("boom")),
+ ]
+
+ executor._create_pods_concurrently()
+
+ executor.kube_scheduler.run_next_batch.assert_called_once()
+ assert executor.task_queue.empty()
+ assert executor.event_buffer[bad_key][0] == State.FAILED
+ assert ok_key not in executor.event_buffer
+
+ @pytest.mark.db_test
+ @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="ExecuteTask workloads require Airflow 3.0+")
+ @pytest.mark.skipif(
+ AirflowKubernetesScheduler is None, reason="kubernetes python package is not installed"
+ )
+ @mock.patch(
+ "airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+ new_callable=mock.AsyncMock,
+ )
+ @mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client")
+ def test_async_pod_creation_with_executetask_workload(
+ self,
+ mock_get_kube_client,
+ mock_kubernetes_job_watcher,
+ mock_get_async_client,
+ create_task_instance,
+ data_file,
+ ):
+ """The concurrent path builds and creates a pod from an AF3 ExecuteTask workload.
+
+ Real AF3 schedulers feed the executor ExecuteTask workloads (command == [workload]),
+ not the legacy ["airflow", "tasks", "run", ...] list. This exercises that branch of
+ the pod build through the concurrent creation path.
+ """
+ from airflow.executors import workloads
+ from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types import KubernetesJob
+
+ template_file = data_file("pods/generator_base_with_secrets.yaml").as_posix()
+ mock_kube_client = mock.MagicMock()
+ mock_kube_client.api_client.sanitize_for_serialization.return_value = {}
+ mock_get_kube_client.return_value = mock_kube_client
+
+ mock_async_api = mock.MagicMock()
+ mock_async_api.create_namespaced_pod = mock.AsyncMock()
+ mock_async_api.api_client.close = mock.AsyncMock()
+ mock_get_async_client.return_value = mock_async_api
+
+ ti = create_task_instance(dag_id="wl_dag", task_id="wl_task", run_id="wl_run")
+ workload = workloads.ExecuteTask.make(ti)
+
+ config = {
+ ("kubernetes_executor", "pod_template_file"): template_file,
+ ("kubernetes_executor", "async_pod_creation"): "True",
+ }
+ with conf_vars(config):
+ executor = KubernetesExecutor()
+ executor.job_id = 5
+ executor._last_completed_pod_adoption = time.monotonic()
+ executor.start()
+ try:
+ job = KubernetesJob(ti.key, [workload], None, template_file)
+ results = executor.kube_scheduler.run_next_batch([job])
+ assert len(results) == 1
+ assert results[0][1] is None # workload built + pod created, no error
+ mock_async_api.create_namespaced_pod.assert_awaited_once()
+ finally:
+ executor.end()
+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor.KubeConfig")
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor.KubernetesExecutor.sync")
@mock.patch("airflow.executors.base_executor.BaseExecutor.trigger_tasks")
diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kube_client.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kube_client.py
new file mode 100644
index 00000000000..53363d49d29
--- /dev/null
+++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kube_client.py
@@ -0,0 +1,39 @@
+# 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.
+from __future__ import annotations
+
+from unittest import mock
+
+import pytest
+
+from airflow.providers.cncf.kubernetes.kube_client import _TimeoutAsyncK8sApiClient, get_async_kube_client
+
+from tests_common.test_utils.config import conf_vars
+
+
+class TestGetAsyncKubeClient:
+ @pytest.mark.asyncio
+ @mock.patch("kubernetes_asyncio.config.load_incluster_config")
+ async def test_wraps_client_with_request_timeout(self, mock_load_incluster):
+ """The async client carries the shared client-side request-timeout wrapper."""
+ with conf_vars(
+ {("kubernetes_executor", "verify_ssl"): "True", ("kubernetes_executor", "ssl_ca_cert"): ""}
+ ):
+ api = await get_async_kube_client(in_cluster=True)
+
+ assert isinstance(api.api_client, _TimeoutAsyncK8sApiClient)
+ mock_load_incluster.assert_called_once()
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 a561db2d7ae..4a641851991 100644
--- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
+++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
@@ -740,6 +740,19 @@ metrics:
legacy_name: "-"
name_variables: []
+ - name: "kubernetes_executor.pod_creation_batch_duration"
+ description: "Milliseconds taken to create one batch of worker pods in a Kubernetes Executor scheduler
+ loop, covering both the sequential and the concurrent (``async_pod_creation``) creation paths."
+ type: "timer"
+ legacy_name: "-"
+ name_variables: []
+
+ - name: "kubernetes_executor.pod_creation_batch_size"
+ description: "Number of worker pods created in one Kubernetes Executor scheduler loop."
+ type: "gauge"
+ legacy_name: "-"
+ name_variables: []
+
- name: "kubernetes_executor.pod_deletion"
description: "Milliseconds taken for a Kubernetes delete_namespaced_pod call from the Kubernetes Executor"
type: "timer"