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 <[email protected]>
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"
