This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 b19bc48f96f Let the EKS executor run per team and next to
KubernetesExecutor (#75449)
b19bc48f96f is described below
commit b19bc48f96f0ae81b5a1b424be5b6efc8dacf55a
Author: Sean Ghaeli <[email protected]>
AuthorDate: Mon Oct 26 14:12:02 2026 -0700
Let the EKS executor run per team and next to KubernetesExecutor (#75449)
The executor installed its client through the [kubernetes_executor]
client_factory setting, so a KubernetesExecutor in the same scope sent its
pods to the EKS cluster too, and it had to refuse that setup. It also
read [aws_eks_executor] only from the global section, so every team ran
on the same cluster, and it refused team values rather than ignore them.
---
providers/amazon/docs/changelog.rst | 6 +
providers/amazon/docs/executors/eks-executor.rst | 52 ++--
providers/amazon/pyproject.toml | 2 +-
.../amazon/aws/executors/eks/_client_factory.py | 30 +-
.../amazon/aws/executors/eks/eks_executor.py | 106 +------
.../providers/amazon/aws/executors/eks/utils.py | 5 +
.../src/airflow/providers/amazon/version_compat.py | 2 -
.../aws/executors/eks/test__client_factory.py | 29 ++
.../amazon/aws/executors/eks/test_eks_executor.py | 330 +++++++++++++--------
9 files changed, 308 insertions(+), 254 deletions(-)
diff --git a/providers/amazon/docs/changelog.rst
b/providers/amazon/docs/changelog.rst
index b0999532227..0b4304c6cc1 100644
--- a/providers/amazon/docs/changelog.rst
+++ b/providers/amazon/docs/changelog.rst
@@ -26,6 +26,12 @@
Changelog
---------
+.. note::
+
+ In multi-team mode, a team's ``AwsEksExecutor`` now reads
``[aws_eks_executor]`` from the team's
+ section and no longer falls back to the global one. Set ``cluster_name``
for each team, and
+ ``region_name`` and ``conn_id`` too if the team needs them.
+
9.39.0
......
diff --git a/providers/amazon/docs/executors/eks-executor.rst
b/providers/amazon/docs/executors/eks-executor.rst
index 12ca2a30ccb..d996f9e38a1 100644
--- a/providers/amazon/docs/executors/eks-executor.rst
+++ b/providers/amazon/docs/executors/eks-executor.rst
@@ -38,17 +38,14 @@ For a quick start guide please see :ref:`here
<eks_setup_guide>`.
Requirements
------------
-The executor needs ``apache-airflow-providers-cncf-kubernetes`` version
10.24.0 or
-newer, the release that added the
-:ref:`apache-airflow-providers-cncf-kubernetes:config:kubernetes_executor__client_factory`
-setting the executor is built on. The executor checks this when it is imported
and raises a
-clear error if an older version is installed. Install the ``cncf.kubernetes``
extra of the
-Amazon provider together with that version:
+The executor needs a version of ``apache-airflow-providers-cncf-kubernetes``
that lets an
+executor built on the Kubernetes executor bring its own Kubernetes client. The
executor checks
+this when it is imported and raises a clear error if an older version is
installed. Install the
+``cncf.kubernetes`` extra of the Amazon provider, which requires that version:
.. code-block:: bash
- pip install 'apache-airflow-providers-amazon[cncf.kubernetes]' \
- 'apache-airflow-providers-cncf-kubernetes>=10.24.0'
+ pip install 'apache-airflow-providers-amazon[cncf.kubernetes]'
The executor also relies on botocore using the regional STS endpoint by
default (botocore 1.40
and newer), and the Amazon provider already requires ``botocore>=1.42.79``, so
no extra install
@@ -121,10 +118,10 @@ See
:doc:`apache-airflow-providers-cncf-kubernetes:kubernetes_executor` for the
for how :ref:`pod templates
<apache-airflow-providers-cncf-kubernetes:concepts:pod_template_file>` and
:ref:`apache-airflow-providers-cncf-kubernetes:concepts:pod_override` work.
-Leave
:ref:`apache-airflow-providers-cncf-kubernetes:config:kubernetes_executor__client_factory`
and
-:ref:`apache-airflow-providers-cncf-kubernetes:config:kubernetes_executor__async_client_factory`
-unset. The executor sets both itself, and raises an error at startup if either
is set to
-something else, so that a conflicting setting cannot silently send tasks to
the wrong cluster.
+The executor builds its own Kubernetes clients and ignores
+:ref:`apache-airflow-providers-cncf-kubernetes:config:kubernetes_executor__client_factory`
and
+:ref:`apache-airflow-providers-cncf-kubernetes:config:kubernetes_executor__async_client_factory`,
+which only apply to the Kubernetes executor.
Set :ref:`[kubernetes_executor] pod_template_file
<apache-airflow-providers-cncf-kubernetes:config:kubernetes_executor__pod_template_file>`
@@ -203,14 +200,29 @@ Multi-team deployments
----------------------
The executor supports :doc:`multi-team
<apache-airflow:core-concepts/multi-team>` mode
-the same way the Kubernetes executor does. Teams share the cluster set in the
global
-:ref:`[aws_eks_executor] <eks_config_options>` section, and each team gets its
own
-team-scoped
:ref:`apache-airflow-providers-cncf-kubernetes:config:kubernetes_executor`
settings,
-such as the namespace and pod template.
-
-The executor refuses to start next to a Kubernetes executor in the same team,
or when both are
-global, because that Kubernetes executor would also send its pods to the EKS
cluster; configure
-them for different teams instead (see `#74844
<https://github.com/apache/airflow/issues/74844>`__).
+the same way the other AWS executors do. Each team reads its own team-scoped
+:ref:`[aws_eks_executor] <eks_config_options>` and
+:ref:`apache-airflow-providers-cncf-kubernetes:config:kubernetes_executor`
settings, so a
+team can use its own cluster, connection, namespace and pod template.
Team-scoped settings
+don't fall back to the global ones, so set ``cluster_name`` for every team
that uses the
+executor.
+
+.. code-block:: ini
+
+ [core]
+ multi_team = True
+ executor =
LocalExecutor;team_a=airflow.providers.amazon.aws.executors.eks.eks_executor.AwsEksExecutor;team_b=airflow.providers.amazon.aws.executors.eks.eks_executor.AwsEksExecutor
+
+ [team_a=aws_eks_executor]
+ cluster_name = team-a-cluster
+
+ [team_b=aws_eks_executor]
+ cluster_name = team-b-cluster
+ conn_id = aws_team_b
+
+The executor can also run alongside the Kubernetes executor, in the same team
or not, and each
+keeps its own cluster. Keep them on different clusters, though: two executors
watching the same
+namespace of one cluster each act on the other's worker pods.
Fault tolerance
---------------
diff --git a/providers/amazon/pyproject.toml b/providers/amazon/pyproject.toml
index d0149c29d47..9ecac3334fa 100644
--- a/providers/amazon/pyproject.toml
+++ b/providers/amazon/pyproject.toml
@@ -91,7 +91,7 @@ dependencies = [
"aiobotocore>=3.0.0",
]
"cncf.kubernetes" = [
- "apache-airflow-providers-cncf-kubernetes>=7.2.0",
+ "apache-airflow-providers-cncf-kubernetes>=7.2.0", # use next version
]
"s3fs" = [
"s3fs>=2023.10.0",
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/_client_factory.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/_client_factory.py
index 4b0dfc2f001..eb0acccf8b8 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/_client_factory.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/_client_factory.py
@@ -16,8 +16,8 @@
# under the License.
"""Kubernetes client factories for the AwsEksExecutor."""
-# Internal to AwsEksExecutor, which writes these functions' import paths into
-# [kubernetes_executor] client_factory and async_client_factory.
+# Internal to AwsEksExecutor, which declares these functions' import paths as
its client_factory
+# and async_client_factory.
from __future__ import annotations
@@ -37,6 +37,7 @@ from airflow.providers.amazon.aws.executors.eks.utils import (
CONFIG_DEFAULTS,
CONFIG_GROUP_NAME,
AllEksConfigKeys,
+ get_cluster_name_env_var,
)
from airflow.providers.amazon.aws.hooks.eks import EksHook
from airflow.providers.amazon.aws.hooks.sts import StsHook
@@ -62,39 +63,46 @@ def _get_sync_token_key(kubernetes_version: str) -> str:
_SYNC_TOKEN_KEY = _get_sync_token_key(kubernetes.__version__)
-def _get_eks_kube_client() -> client.CoreV1Api:
- """Build a Kubernetes client for the configured EKS cluster, in memory and
without a kubeconfig."""
+def _get_eks_kube_client(team_name: str | None = None) -> client.CoreV1Api:
+ """Build a Kubernetes client for the team's EKS cluster, in memory and
without a kubeconfig."""
from kubernetes import client
configuration = client.Configuration()
- _configure_eks_auth(configuration, _SYNC_TOKEN_KEY)
+ _configure_eks_auth(configuration, _SYNC_TOKEN_KEY, team_name)
return client.CoreV1Api(client.ApiClient(configuration=configuration))
-def _get_eks_async_kube_client() -> async_client.CoreV1Api:
+def _get_eks_async_kube_client(team_name: str | None = None) ->
async_client.CoreV1Api:
"""Build the asynchronous Kubernetes client used when
``async_pod_creation`` is enabled."""
from kubernetes_asyncio import client as async_client
configuration = async_client.Configuration()
# kubernetes_asyncio has always used "BearerToken".
- _configure_eks_auth(configuration, "BearerToken")
+ _configure_eks_auth(configuration, "BearerToken", team_name)
return
async_client.CoreV1Api(async_client.ApiClient(configuration=configuration))
def _configure_eks_auth(
- configuration: client.Configuration | async_client.Configuration,
token_key: str
+ configuration: client.Configuration | async_client.Configuration,
+ token_key: str,
+ team_name: str | None,
) -> None:
- cluster_name = conf.get(CONFIG_GROUP_NAME, AllEksConfigKeys.CLUSTER_NAME,
fallback=None)
+ # The factories run where there is no executor, such as the pod watcher
process, so they look
+ # up the team's config themselves, the same way ExecutorConf does. Airflow
releases without
+ # multi-team don't accept team_name.
+ team_kwargs = {"team_name": team_name} if team_name else {}
+ cluster_name = conf.get(CONFIG_GROUP_NAME, AllEksConfigKeys.CLUSTER_NAME,
fallback=None, **team_kwargs)
if not cluster_name:
raise ValueError(
f"[{CONFIG_GROUP_NAME}] cluster_name is required to build an EKS
client. "
- "Set it in airflow.cfg or with
AIRFLOW__AWS_EKS_EXECUTOR__CLUSTER_NAME."
+ f"Set it in airflow.cfg or with
{get_cluster_name_env_var(team_name)}."
)
- region_name = conf.get(CONFIG_GROUP_NAME, AllEksConfigKeys.REGION_NAME,
fallback=None)
+ region_name = conf.get(CONFIG_GROUP_NAME, AllEksConfigKeys.REGION_NAME,
fallback=None, **team_kwargs)
conn_id = conf.get(
CONFIG_GROUP_NAME,
AllEksConfigKeys.AWS_CONN_ID,
fallback=CONFIG_DEFAULTS[AllEksConfigKeys.AWS_CONN_ID],
+ **team_kwargs,
)
eks_hook = EksHook(aws_conn_id=conn_id, region_name=region_name)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/eks_executor.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/eks_executor.py
index 083b0bb9d8d..281fe405f14 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/eks_executor.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/eks_executor.py
@@ -18,41 +18,34 @@
from __future__ import annotations
-import os
from typing import TYPE_CHECKING
-from airflow.executors.executor_loader import ExecutorLoader
from airflow.providers.amazon.aws.executors.eks.utils import (
CONFIG_DEFAULTS,
CONFIG_GROUP_NAME,
AllEksConfigKeys,
+ get_cluster_name_env_var,
)
-from airflow.providers.amazon.version_compat import AIRFLOW_V_3_2_PLUS
-from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, conf
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
-_CNCF_KUBERNETES_REQUIRED = "AwsEksExecutor requires
apache-airflow-providers-cncf-kubernetes>=10.24.0"
+_CNCF_KUBERNETES_REQUIRED = (
+ "AwsEksExecutor requires the apache-airflow-providers-cncf-kubernetes
version that the "
+ "Amazon provider's cncf.kubernetes extra installs: "
+ "pip install 'apache-airflow-providers-amazon[cncf.kubernetes]'"
+)
try:
from kubernetes.client.rest import ApiException
from airflow.providers.cncf.kubernetes.executors.kubernetes_executor
import KubernetesExecutor
- from airflow.providers.cncf.kubernetes.get_provider_info import (
- get_provider_info as get_cncf_kubernetes_provider_info,
- )
except ImportError as e:
raise AirflowOptionalProviderFeatureException(_CNCF_KUBERNETES_REQUIRED)
from e
-# Check for the option, not the version, since unreleased trees report the
previous version.
-if "client_factory" not in (
- get_cncf_kubernetes_provider_info().get("config",
{}).get("kubernetes_executor", {}).get("options", {})
-):
+# Check for the feature, not the version, since unreleased trees report the
previous version.
+if not hasattr(KubernetesExecutor, "client_factory"):
raise AirflowOptionalProviderFeatureException(_CNCF_KUBERNETES_REQUIRED)
_FACTORY_MODULE = "airflow.providers.amazon.aws.executors.eks._client_factory"
-_CLIENT_FACTORY_PATH = f"{_FACTORY_MODULE}._get_eks_kube_client"
-_ASYNC_CLIENT_FACTORY_PATH = f"{_FACTORY_MODULE}._get_eks_async_kube_client"
-# Airflow 2 hybrid executors that wrap a KubernetesExecutor.
-_HYBRID_KUBERNETES_EXECUTORS = ("CeleryKubernetesExecutor",
"LocalKubernetesExecutor")
class AwsEksExecutor(KubernetesExecutor):
@@ -65,19 +58,18 @@ class AwsEksExecutor(KubernetesExecutor):
"""
supports_multi_team: bool = True
+ client_factory = f"{_FACTORY_MODULE}._get_eks_kube_client"
+ async_client_factory = f"{_FACTORY_MODULE}._get_eks_async_kube_client"
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
- # These need team_name, which super().__init__ sets.
- self._check_no_team_eks_config()
+ # Needs team_name, which super().__init__ sets.
self._validate_eks_config()
- self._check_no_kubernetes_executor_in_scope()
- self._ensure_client_factory()
def start(self) -> None:
"""Call this when the Executor is run for the first time by the
scheduler."""
super().start()
- check_health = conf.getboolean(
+ check_health = self.conf.getboolean(
CONFIG_GROUP_NAME,
AllEksConfigKeys.CHECK_HEALTH_ON_STARTUP,
fallback=CONFIG_DEFAULTS[AllEksConfigKeys.CHECK_HEALTH_ON_STARTUP],
@@ -115,75 +107,9 @@ class AwsEksExecutor(KubernetesExecutor):
) from e
self.log.info("EKS Executor health check has succeeded.")
- @staticmethod
- def _validate_eks_config() -> None:
- if not conf.get(CONFIG_GROUP_NAME, AllEksConfigKeys.CLUSTER_NAME,
fallback=None):
+ def _validate_eks_config(self) -> None:
+ if not self.conf.get(CONFIG_GROUP_NAME, AllEksConfigKeys.CLUSTER_NAME,
fallback=None):
raise ValueError(
f"AwsEksExecutor requires [{CONFIG_GROUP_NAME}] cluster_name
to be set. "
- "Set it in airflow.cfg or with
AIRFLOW__AWS_EKS_EXECUTOR__CLUSTER_NAME."
- )
-
- def _check_no_team_eks_config(self) -> None:
- # [aws_eks_executor] is only read globally for now, so refuse team
values rather than ignore them.
- if not self.team_name:
- return
- # self.conf looks options up for this team only, falling back to the
built-in defaults.
- team_keys = sorted(
- key
- for key in AllEksConfigKeys()
- if self.conf.get(CONFIG_GROUP_NAME, key, fallback=None) not in
(None, CONFIG_DEFAULTS.get(key))
- )
- if team_keys:
- raise ValueError(
- f"AwsEksExecutor doesn't support team-specific
[{CONFIG_GROUP_NAME}] config yet, but team "
- f"{self.team_name} sets {', '.join(team_keys)}. Set these in
the global "
- f"[{CONFIG_GROUP_NAME}] section instead. Per-team clusters are
tracked at "
- "https://github.com/apache/airflow/issues/74844";
+ f"Set it in airflow.cfg or with
{get_cluster_name_env_var(self.team_name)}."
)
-
- def _check_no_kubernetes_executor_in_scope(self) -> None:
- executor_names = (
- ExecutorLoader.get_executor_names(validate_teams=False)
- if AIRFLOW_V_3_2_PLUS
- else ExecutorLoader.get_executor_names()
- )
- for executor_name in executor_names:
- # ExecutorName has no team_name before Airflow 3.1.
- if getattr(executor_name, "team_name", None) != self.team_name:
- continue
- executor_cls, _ = ExecutorLoader.import_executor_cls(executor_name)
- if issubclass(executor_cls, AwsEksExecutor) or not (
- issubclass(executor_cls, KubernetesExecutor)
- or executor_cls.__name__ in _HYBRID_KUBERNETES_EXECUTORS
- ):
- continue
- scope = f"team {self.team_name}" if self.team_name else "the
global scope"
- raise ValueError(
- f"AwsEksExecutor can't run alongside {executor_cls.__name__}
in {scope}. AwsEksExecutor "
- f"sets the [kubernetes_executor] client factory for {scope},
so {executor_cls.__name__} "
- "would silently send its pods to the EKS cluster. Configure
the two executors for "
- "different teams, or don't use them together. "
- "See https://github.com/apache/airflow/issues/74844";
- )
-
- def _ensure_client_factory(self) -> None:
- # A team's factory is looked up in team config only, with no global
fallback.
- team_name = self.team_name
- team_kwargs = {"team_name": team_name} if team_name else {}
- env_var_prefix = f"AIRFLOW__{team_name.upper()}___" if team_name else
"AIRFLOW__"
- for key, path in (
- ("client_factory", _CLIENT_FACTORY_PATH),
- ("async_client_factory", _ASYNC_CLIENT_FACTORY_PATH),
- ):
- configured = conf.get("kubernetes_executor", key, fallback=None,
**team_kwargs)
- if configured and configured != path:
- raise ValueError(
- f"AwsEksExecutor sets [kubernetes_executor] {key} itself,
so leave it unset; "
- f"it is set to {configured}"
- )
- if not configured:
- # An env var, not conf.set, so the spawned pod watcher process
sees it too. A plain
- # KubernetesExecutor in the same scope reads it as well, which
- # _check_no_kubernetes_executor_in_scope refuses; replacing
this env var is tracked at
- # https://github.com/apache/airflow/issues/74844
-
os.environ[f"{env_var_prefix}KUBERNETES_EXECUTOR__{key.upper()}"] = path
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/utils.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/utils.py
index 52e0064f9d8..d97b8162b38 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/utils.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/utils.py
@@ -35,3 +35,8 @@ class AllEksConfigKeys(BaseConfigKeys):
CHECK_HEALTH_ON_STARTUP = "check_health_on_startup"
CLUSTER_NAME = "cluster_name"
REGION_NAME = "region_name"
+
+
+def get_cluster_name_env_var(team_name: str | None) -> str:
+ team_prefix = f"{team_name.upper()}___" if team_name else ""
+ return
f"AIRFLOW__{team_prefix}{CONFIG_GROUP_NAME.upper()}__{AllEksConfigKeys.CLUSTER_NAME.upper()}"
diff --git a/providers/amazon/src/airflow/providers/amazon/version_compat.py
b/providers/amazon/src/airflow/providers/amazon/version_compat.py
index 92e9bc41c2c..fa6faf2ba41 100644
--- a/providers/amazon/src/airflow/providers/amazon/version_compat.py
+++ b/providers/amazon/src/airflow/providers/amazon/version_compat.py
@@ -40,7 +40,6 @@ AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3,
0, 0)
AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0)
AIRFLOW_V_3_1_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 1)
AIRFLOW_V_3_1_8_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 8)
-AIRFLOW_V_3_2_PLUS: bool = get_base_airflow_version_tuple() >= (3, 2, 0)
AIRFLOW_V_3_3_PLUS: bool = get_base_airflow_version_tuple() >= (3, 3, 0)
AIRFLOW_V_3_4_PLUS: bool = get_base_airflow_version_tuple() >= (3, 4, 0)
@@ -61,7 +60,6 @@ __all__ = [
"AIRFLOW_V_3_1_PLUS",
"AIRFLOW_V_3_1_1_PLUS",
"AIRFLOW_V_3_1_8_PLUS",
- "AIRFLOW_V_3_2_PLUS",
"AIRFLOW_V_3_3_PLUS",
"AIRFLOW_V_3_4_PLUS",
"NOTSET",
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/eks/test__client_factory.py
b/providers/amazon/tests/unit/amazon/aws/executors/eks/test__client_factory.py
index 45306e27316..f3b10ad6bac 100644
---
a/providers/amazon/tests/unit/amazon/aws/executors/eks/test__client_factory.py
+++
b/providers/amazon/tests/unit/amazon/aws/executors/eks/test__client_factory.py
@@ -43,6 +43,7 @@ from airflow.providers.amazon.aws.utils.eks_get_token import (
)
from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS
CLUSTER_NAME = "test-eks-cluster"
CLUSTER_ENDPOINT = "https://ABCDEF0123456789.gr7.us-east-1.eks.amazonaws.com";
@@ -194,6 +195,34 @@ class TestGetEksKubeClient:
):
_get_eks_kube_client()
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="Multi-team requires
Airflow 3.2+")
+ def test_team_without_cluster_name_raises(self, mock_aws):
+ """The global cluster_name set by mock_aws is not a fallback for the
team."""
+ with pytest.raises(
+ ValueError,
+ match=r"cluster_name is
required.*AIRFLOW__TEAM_A___AWS_EKS_EXECUTOR__CLUSTER_NAME",
+ ):
+ _get_eks_kube_client(team_name="team_a")
+
+
[email protected](not AIRFLOW_V_3_2_PLUS, reason="Multi-team requires
Airflow 3.2+")
[email protected]("factory", [_get_eks_kube_client,
_get_eks_async_kube_client], ids=["sync", "async"])
[email protected]
+# kubernetes_asyncio loads the CA eagerly and the fake PEM is not loadable.
[email protected](_client_factory, "_write_cluster_ca_file",
return_value=None)
+async def test_builds_client_for_the_teams_cluster(_, mock_aws, monkeypatch,
factory):
+ monkeypatch.setenv("AIRFLOW__TEAM_A___AWS_EKS_EXECUTOR__CLUSTER_NAME",
"team-a-cluster")
+ monkeypatch.setenv("AIRFLOW__TEAM_A___AWS_EKS_EXECUTOR__REGION_NAME",
"eu-west-1")
+ monkeypatch.setenv("AIRFLOW__TEAM_A___AWS_EKS_EXECUTOR__CONN_ID",
"aws_team_a")
+
+ core_v1 = factory(team_name="team_a")
+
+ mock_aws["eks_hook_cls"].assert_called_once_with(aws_conn_id="aws_team_a",
region_name="eu-west-1")
+
mock_aws["eks_hook"].conn.describe_cluster.assert_called_once_with(name="team-a-cluster")
+ assert mock_aws["fetch_token"].call_args.args[0] == "team-a-cluster"
+ if inspect.iscoroutinefunction(core_v1.api_client.close):
+ await core_v1.api_client.close()
+
@pytest.mark.asyncio
# kubernetes_asyncio loads the CA eagerly and the fake PEM is not loadable.
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/eks/test_eks_executor.py
b/providers/amazon/tests/unit/amazon/aws/executors/eks/test_eks_executor.py
index 096dce5c64b..db78f1bab42 100644
--- a/providers/amazon/tests/unit/amazon/aws/executors/eks/test_eks_executor.py
+++ b/providers/amazon/tests/unit/amazon/aws/executors/eks/test_eks_executor.py
@@ -16,10 +16,11 @@
# under the License.
from __future__ import annotations
+import argparse
import importlib
-import os
import sys
-from contextlib import nullcontext
+from base64 import b64encode
+from datetime import UTC, datetime
from unittest import mock
import pytest
@@ -27,22 +28,17 @@ from kubernetes.client import CoreV1Api
from kubernetes.client.rest import ApiException
from airflow.providers.amazon.aws.executors.eks import AwsEksExecutor,
_client_factory
-from airflow.providers.amazon.aws.executors.eks.eks_executor import (
- _ASYNC_CLIENT_FACTORY_PATH,
- _CLIENT_FACTORY_PATH,
-)
+from airflow.providers.cncf.kubernetes.cli import kubernetes_command
from airflow.providers.cncf.kubernetes.executors.kubernetes_executor import
KubernetesExecutor
-from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, conf
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS
EKS_EXECUTOR_MODULE = "airflow.providers.amazon.aws.executors.eks.eks_executor"
-EKS_EXECUTOR_PATH = f"{EKS_EXECUTOR_MODULE}.AwsEksExecutor"
-LOCAL_KUBERNETES_EXECUTOR_PATH = (
-
"airflow.providers.cncf.kubernetes.executors.local_kubernetes_executor.LocalKubernetesExecutor"
-)
-CLIENT_FACTORY_ENV_VAR = "AIRFLOW__KUBERNETES_EXECUTOR__CLIENT_FACTORY"
+CLUSTER_ENDPOINT = "https://ABCDEF0123456789.gr7.us-east-1.eks.amazonaws.com";
+JOB_WATCHER =
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher"
+USER_CLIENT = "client-from-user-factory"
REQUIRES_MULTI_TEAM = pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS,
reason="Multi-team requires Airflow 3.2+")
TEAMS = [
pytest.param(None, "AIRFLOW__", id="global"),
@@ -59,132 +55,166 @@ def _team_kwargs(team_name: str | None) -> dict:
return {"team_name": team_name} if team_name else {}
+def build_user_client():
+ return USER_CLIENT
+
+
@pytest.fixture(autouse=True)
def eks_config():
- with mock.patch.dict(os.environ), conf_vars({("aws_eks_executor",
"cluster_name"): "test-eks-cluster"}):
- for prefix in ("AIRFLOW__", "AIRFLOW__TEAM_A___"):
- for key in ("CLIENT_FACTORY", "ASYNC_CLIENT_FACTORY"):
- os.environ.pop(f"{prefix}KUBERNETES_EXECUTOR__{key}", None)
+ with conf_vars({("aws_eks_executor", "cluster_name"): "test-eks-cluster"}):
yield
-class TestAwsEksExecutor:
- @conf_vars({("aws_eks_executor", "cluster_name"): None})
- def test_missing_cluster_name_raises(self):
- with pytest.raises(
- ValueError, match=r"\[aws_eks_executor\]
cluster_name.*AIRFLOW__AWS_EKS_EXECUTOR__CLUSTER_NAME"
- ):
- AwsEksExecutor()
-
- @pytest.mark.parametrize(("team_name", "env_prefix"), TEAMS)
- def test_sets_client_factories_when_unset(self, team_name, env_prefix):
- AwsEksExecutor(**_team_kwargs(team_name))
-
- assert {k: v for k, v in os.environ.items() if
k.endswith("CLIENT_FACTORY")} == {
- f"{env_prefix}KUBERNETES_EXECUTOR__CLIENT_FACTORY":
_CLIENT_FACTORY_PATH,
- f"{env_prefix}KUBERNETES_EXECUTOR__ASYNC_CLIENT_FACTORY":
_ASYNC_CLIENT_FACTORY_PATH,
[email protected]
+def mock_aws():
+ with (
+ mock.patch.object(_client_factory, "EksHook") as eks_hook_cls,
+ mock.patch.object(_client_factory, "StsHook") as sts_hook_cls,
+ mock.patch.object(_client_factory, "fetch_access_token_for_cluster",
autospec=True) as fetch_token,
+ ):
+ eks_hook = eks_hook_cls.return_value
+ eks_hook.conn.describe_cluster.return_value = {
+ "cluster": {
+ "status": "ACTIVE",
+ "endpoint": CLUSTER_ENDPOINT,
+ "certificateAuthority": {"data":
b64encode(b"fake-ca").decode()},
+ }
}
- team_kwargs = _team_kwargs(team_name)
- assert (
- conf.getimport("kubernetes_executor", "client_factory",
**team_kwargs)
- is _client_factory._get_eks_kube_client
- )
- assert (
- conf.getimport("kubernetes_executor", "async_client_factory",
**team_kwargs)
- is _client_factory._get_eks_async_kube_client
- )
+ eks_hook.get_session.return_value.region_name = "us-east-1"
+ sts_hook_cls.return_value.conn_client_meta.endpoint_url =
"https://sts.us-east-1.amazonaws.com";
+ fetch_token.return_value = "k8s-aws-v1.token"
+ yield {"eks_hook_cls": eks_hook_cls, "eks_hook": eks_hook,
"fetch_token": fetch_token}
- @conf_vars({("kubernetes_executor", "client_factory"):
_CLIENT_FACTORY_PATH})
- def test_accepts_already_matching_client_factory(self):
- AwsEksExecutor()
- assert CLIENT_FACTORY_ENV_VAR not in os.environ
+def _start_and_end(executor: KubernetesExecutor) -> None:
+ executor.job_id = 1
+ try:
+ executor.start()
+ finally:
+ executor.end()
- @pytest.mark.parametrize("key", ["client_factory", "async_client_factory"])
- @pytest.mark.parametrize(("team_name", "env_prefix"), TEAMS)
- def test_conflicting_client_factory_raises(self, monkeypatch, team_name,
env_prefix, key):
- monkeypatch.setenv(f"{env_prefix}KUBERNETES_EXECUTOR__{key.upper()}",
"my_company.build_client")
- with pytest.raises(ValueError, match=rf"{key}
itself.*my_company.build_client"):
- AwsEksExecutor(**_team_kwargs(team_name))
+class TestAwsEksExecutor:
+ @pytest.mark.parametrize(("team_name", "env_prefix"), TEAMS)
+ @conf_vars({("aws_eks_executor", "cluster_name"): None})
+ def test_missing_cluster_name_raises(self, monkeypatch, team_name,
env_prefix):
+ if team_name:
+ # Team config doesn't fall back to the global section.
+ monkeypatch.setenv("AIRFLOW__AWS_EKS_EXECUTOR__CLUSTER_NAME",
"global-cluster")
- @REQUIRES_MULTI_TEAM
- @pytest.mark.parametrize(
- ("env", "file_config", "key"),
- [
- pytest.param(
- {"AIRFLOW__TEAM_A___AWS_EKS_EXECUTOR__CLUSTER_NAME":
"team-cluster"},
- {("aws_eks_executor", "cluster_name"): None},
- "cluster_name",
- id="env-var-without-global-cluster",
- ),
- pytest.param(
- {"AIRFLOW__TEAM_A___AWS_EKS_EXECUTOR__CONN_ID": "team_aws"},
- {},
- "conn_id",
- id="env-var-for-option-with-default",
- ),
- pytest.param(
- {},
- {("team_a=aws_eks_executor", "region_name"): "eu-west-1"},
- "region_name",
- id="config-file-section",
- ),
- ],
- )
- def test_team_eks_config_raises(self, monkeypatch, env, file_config, key):
- for name, value in env.items():
- monkeypatch.setenv(name, value)
-
- with (
- conf_vars(file_config),
- pytest.raises(
- ValueError, match=rf"team-specific \[aws_eks_executor\].*team
team_a sets {key}\..*74844"
- ),
+ with pytest.raises(
+ ValueError,
+ match=rf"\[aws_eks_executor\]
cluster_name.*{env_prefix}AWS_EKS_EXECUTOR__CLUSTER_NAME",
):
- AwsEksExecutor(team_name="team_a")
+ AwsEksExecutor(**_team_kwargs(team_name))
@pytest.mark.parametrize(
- ("executors", "team_name", "conflict"),
+ ("team_name", "cluster_name", "conn_id", "region_name",
"check_health"),
[
pytest.param(
- f"KubernetesExecutor,{EKS_EXECUTOR_PATH}",
- None,
- "KubernetesExecutor in the global scope",
- id="both-global",
- ),
- pytest.param(
- f"{LOCAL_KUBERNETES_EXECUTOR_PATH},{EKS_EXECUTOR_PATH}",
- None,
- "LocalKubernetesExecutor in the global scope",
- id="hybrid-global",
- ),
- pytest.param(
- f"LocalExecutor;team_a=KubernetesExecutor,{EKS_EXECUTOR_PATH}",
"team_a",
- "KubernetesExecutor in team team_a",
- id="both-same-team",
+ "team-a-cluster",
+ "aws_team_a",
+ "eu-west-1",
+ True,
+ id="team",
marks=REQUIRES_MULTI_TEAM,
),
- pytest.param(
- f"KubernetesExecutor;team_a={EKS_EXECUTOR_PATH}",
- "team_a",
- None,
- id="global-kubernetes-team-eks",
- marks=REQUIRES_MULTI_TEAM,
- ),
- pytest.param(EKS_EXECUTOR_PATH, None, None, id="eks-alone"),
+ pytest.param(None, "global-cluster", "aws_global", "us-east-1",
False, id="global"),
],
)
- @pytest.mark.usefixtures("clean_executor_loader")
- def test_refuses_kubernetes_executor_in_same_scope(self, executors,
team_name, conflict):
- expectation = (
- pytest.raises(ValueError, match=rf"alongside
{conflict}\..*issues/74844")
- if conflict
- else nullcontext()
+ # kubernetes_asyncio loads the CA eagerly and the fake one is not loadable.
+ @mock.patch.object(_client_factory, "_write_cluster_ca_file",
return_value=None)
+ @mock.patch.object(AwsEksExecutor, "check_health", autospec=True)
+ @mock.patch(JOB_WATCHER, autospec=True)
+ def test_team_config(
+ self,
+ mock_watcher,
+ mock_check_health,
+ _,
+ mock_aws,
+ monkeypatch,
+ team_name,
+ cluster_name,
+ conn_id,
+ region_name,
+ check_health,
+ ):
+ """Each executor builds both its clients from its own
[aws_eks_executor] values, and so does its watcher."""
+ for key, value in {
+ "AIRFLOW__KUBERNETES_EXECUTOR__ASYNC_POD_CREATION": "True",
+ "AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__ASYNC_POD_CREATION":
"True",
+ "AIRFLOW__AWS_EKS_EXECUTOR__CLUSTER_NAME": "global-cluster",
+ "AIRFLOW__AWS_EKS_EXECUTOR__REGION_NAME": "us-east-1",
+ "AIRFLOW__AWS_EKS_EXECUTOR__CONN_ID": "aws_global",
+ "AIRFLOW__AWS_EKS_EXECUTOR__CHECK_HEALTH_ON_STARTUP": "False",
+ # The team leaves check_health_on_startup at its default rather
than the global value.
+ "AIRFLOW__TEAM_A___AWS_EKS_EXECUTOR__CLUSTER_NAME":
"team-a-cluster",
+ "AIRFLOW__TEAM_A___AWS_EKS_EXECUTOR__REGION_NAME": "eu-west-1",
+ "AIRFLOW__TEAM_A___AWS_EKS_EXECUTOR__CONN_ID": "aws_team_a",
+ }.items():
+ monkeypatch.setenv(key, value)
+ executor = AwsEksExecutor(**_team_kwargs(team_name))
+ executor.job_id = 1
+ executor.start()
+ try:
+ # Builds the async client, as the first batch of pods would.
+ executor.kube_scheduler._run_pods_async([])
+ async_client = executor.kube_scheduler._async_pod_client
+ finally:
+ executor.end()
+
+ assert executor.kube_client.api_client.configuration.host ==
CLUSTER_ENDPOINT
+ assert async_client.api_client.configuration.host == CLUSTER_ENDPOINT
+ assert (
+ mock_aws["eks_hook_cls"].call_args_list
+ == [mock.call(aws_conn_id=conn_id, region_name=region_name)] * 2
)
- with conf_vars({("core", "executor"): executors, ("core",
"multi_team"): "True"}), expectation:
- AwsEksExecutor(**_team_kwargs(team_name))
+ assert mock_aws["eks_hook"].conn.describe_cluster.call_args_list ==
[mock.call(name=cluster_name)] * 2
+ assert {call.args[0] for call in
mock_aws["fetch_token"].call_args_list} == {cluster_name}
+ assert mock_check_health.call_count == int(check_health)
+ assert mock_watcher.call_args.kwargs["team_name"] == team_name
+ assert mock_watcher.call_args.kwargs["client_factory"] ==
AwsEksExecutor.client_factory
+
+ @pytest.mark.parametrize("user_factory", [False, True],
ids=["default-client", "user-factory"])
+ @pytest.mark.parametrize(("team_name", "env_prefix"), TEAMS)
+ @mock.patch.object(AwsEksExecutor, "check_health", autospec=True)
+ @mock.patch(JOB_WATCHER, autospec=True)
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client",
autospec=True)
+ def test_runs_alongside_kubernetes_executor(
+ self,
+ mock_get_kube_client,
+ mock_watcher,
+ mock_check_health,
+ mock_aws,
+ monkeypatch,
+ team_name,
+ env_prefix,
+ user_factory,
+ ):
+ """A KubernetesExecutor in the same scope keeps its own client,
including its own factory."""
+ monkeypatch.setenv(f"{env_prefix}AWS_EKS_EXECUTOR__CLUSTER_NAME",
"test-eks-cluster")
+ if user_factory:
+ monkeypatch.setenv(
+ f"{env_prefix}KUBERNETES_EXECUTOR__CLIENT_FACTORY",
f"{__name__}.build_user_client"
+ )
+ # The scheduler builds every executor before it starts any of them.
+ eks_executor = AwsEksExecutor(**_team_kwargs(team_name))
+ kubernetes_executor = KubernetesExecutor(**_team_kwargs(team_name))
+
+ _start_and_end(eks_executor)
+ _start_and_end(kubernetes_executor)
+
+ assert eks_executor.kube_client.api_client.configuration.host ==
CLUSTER_ENDPOINT
+ if user_factory:
+ assert kubernetes_executor.kube_client == USER_CLIENT
+ else:
+ assert kubernetes_executor.kube_client is
mock_get_kube_client.return_value
+ mock_get_kube_client.assert_called_once_with(team_name=team_name)
+ assert [call.kwargs["client_factory"] for call in
mock_watcher.call_args_list] == [
+ AwsEksExecutor.client_factory,
+ None,
+ ]
@pytest.mark.parametrize("check_health", [True, False])
@mock.patch.object(KubernetesExecutor, "start", autospec=True)
@@ -217,22 +247,62 @@ class TestCncfKubernetesRequirement:
sys.modules,
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor", None
)
- with pytest.raises(AirflowOptionalProviderFeatureException,
match=">=10.24.0"):
+ with pytest.raises(
+ AirflowOptionalProviderFeatureException,
+ match=r"apache-airflow-providers-amazon\[cncf.kubernetes\]",
+ ):
importlib.import_module(EKS_EXECUTOR_MODULE)
- @pytest.mark.parametrize(
- "provider_info",
- [
- pytest.param({"config": {"kubernetes_executor": {"options": {}}}},
id="no-client-factory"),
- pytest.param({}, id="no-config"),
- ],
- )
-
@mock.patch("airflow.providers.cncf.kubernetes.get_provider_info.get_provider_info",
autospec=True)
- def test_import_fails_with_cncf_kubernetes_without_client_factory(
- self, mock_get_provider_info, monkeypatch, provider_info
- ):
- mock_get_provider_info.return_value = provider_info
+ def
test_import_fails_with_cncf_kubernetes_without_executor_client_factory(self,
monkeypatch):
+ monkeypatch.delattr(KubernetesExecutor, "client_factory")
monkeypatch.delitem(sys.modules, EKS_EXECUTOR_MODULE)
- with pytest.raises(AirflowOptionalProviderFeatureException,
match=">=10.24.0"):
+ with pytest.raises(
+ AirflowOptionalProviderFeatureException,
+ match=r"apache-airflow-providers-amazon\[cncf.kubernetes\]",
+ ):
importlib.import_module(EKS_EXECUTOR_MODULE)
+
+
[email protected]_test
+class TestCleanupPods:
+ @REQUIRES_MULTI_TEAM
+ @pytest.mark.usefixtures("clean_executor_loader")
+ @conf_vars(
+ {
+ ("core", "multi_team"): "True",
+ ("core", "executor"):
f"LocalExecutor;team_a={EKS_EXECUTOR_MODULE}.AwsEksExecutor",
+ }
+ )
+ @mock.patch.object(kubernetes_command, "_delete_pod", autospec=True)
+ @mock.patch.object(CoreV1Api, "list_namespaced_pod", autospec=True)
+ def test_cleanup_pods_team_uses_the_teams_eks_cluster(
+ self, mock_list_pods, mock_delete_pod, mock_aws, monkeypatch
+ ):
+ """``airflow kubernetes cleanup-pods --team`` connects through the
team's AwsEksExecutor."""
+ monkeypatch.setenv("AIRFLOW__TEAM_A___AWS_EKS_EXECUTOR__CLUSTER_NAME",
"team-a-cluster")
+ monkeypatch.setenv("AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__NAMESPACE",
"team-a-namespace")
+ pod = mock.MagicMock()
+ pod.metadata.name = "finished-pod"
+ pod.metadata.creation_timestamp = datetime(2021, 12, 20, tzinfo=UTC)
+ pod.status.phase = "Succeeded"
+ pod.status.reason = None
+ mock_list_pods.return_value = mock.MagicMock(items=[pod],
metadata=mock.MagicMock(_continue=None))
+
+ kubernetes_command.cleanup_pods(
+ argparse.Namespace(
+ namespace=None,
+ min_pending_minutes=30,
+ min_completed_minutes=1,
+ team="team_a",
+ executor=None,
+ allow_shared_namespace=False,
+ verbose=False,
+ )
+ )
+
+
mock_aws["eks_hook"].conn.describe_cluster.assert_called_once_with(name="team-a-cluster")
+ kube_client = mock_list_pods.call_args.args[0]
+ assert kube_client.api_client.configuration.host == CLUSTER_ENDPOINT
+ assert mock_list_pods.call_args.kwargs["namespace"] ==
"team-a-namespace"
+ mock_delete_pod.assert_called_once_with("finished-pod",
"team-a-namespace", kube_client)
