ferruzzi commented on code in PR #74399: URL: https://github.com/apache/airflow/pull/74399#discussion_r4288802873
########## providers/amazon/src/airflow/providers/amazon/aws/executors/eks/eks_executor.py: ########## @@ -0,0 +1,137 @@ +# 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. +"""AwsEksExecutor: run Airflow tasks as pods on an Amazon EKS cluster.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +from airflow.providers.amazon.aws.executors.eks.utils import ( + CONFIG_DEFAULTS, + CONFIG_GROUP_NAME, + AllEksConfigKeys, +) +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, conf + +_CNCF_KUBERNETES_REQUIRED = "AwsEksExecutor requires apache-airflow-providers-cncf-kubernetes>=10.24.0" + +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", {}) +): + 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" + + +class AwsEksExecutor(KubernetesExecutor): + """ + A KubernetesExecutor that authenticates against an Amazon EKS cluster. + + Builds the Kubernetes client from ``[aws_eks_executor]`` configuration and keeps the + short-lived EKS token fresh in-process. All pod-level behaviour comes unchanged from the + KubernetesExecutor and its ``[kubernetes_executor]`` configuration. + """ + + supports_multi_team: bool = True + + def __init__(self, *args, **kwargs): + self._validate_eks_config() + super().__init__(*args, **kwargs) + # Needs team_name, which super().__init__ sets. + 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( + CONFIG_GROUP_NAME, + AllEksConfigKeys.CHECK_HEALTH_ON_STARTUP, + fallback=CONFIG_DEFAULTS[AllEksConfigKeys.CHECK_HEALTH_ON_STARTUP], + ) + + if not check_health: + return + + self.log.info("Starting EKS Executor and determining health...") + try: + self.check_health() + except RuntimeError: + self.log.error("Stopping the Airflow Scheduler from starting until the issue is resolved.") + raise + + def check_health(self) -> None: + """ + Make a test Kubernetes API call to check the health of the EKS Executor. + + Building the client only proves the cluster exists; without this, a missing EKS access + entry or RBAC binding only shows up later as a watcher error loop. + """ + if TYPE_CHECKING: + assert self.kube_client + namespace = self.kube_config.kube_namespace + try: + self.kube_client.list_namespaced_pod(namespace, limit=1) + except ApiException as e: + raise RuntimeError( + f"EKS Executor health check has failed because: cannot list pods in namespace {namespace} " + f"({e.status} {e.reason}). The IAM principal behind the AWS connection needs an EKS access " + "entry (or aws-auth mapping) on the cluster that allows create, get, list, watch, patch and " + f"delete on pods and get on pods/log in namespace {namespace}. See 'Grant access to the " + "cluster' in the AwsEksExecutor docs." + ) 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): Review Comment: This is maybe a @vincbeck multi-team config question... ECS, Batch, and Lambda executors all handle `conf.get` calls through `self.conf` as part of the multi-team support. Your docs say this new executor supports multi-team but you are not following that same pattern here and elsewhere. Is that a miss, or is there a reason EKS doesn't need to match the same pattern? ########## providers/amazon/src/airflow/providers/amazon/aws/executors/eks/_client_factory.py: ########## @@ -0,0 +1,144 @@ +# 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. +"""Kubernetes client factories for the AwsEksExecutor.""" + +# Internal to AwsEksExecutor, which writes these functions' import paths into +# [kubernetes_executor] client_factory and async_client_factory. + +from __future__ import annotations + +import os +import tempfile +import time +from base64 import b64decode +from typing import TYPE_CHECKING + +import kubernetes +from packaging.version import Version + +from airflow.providers.amazon.aws.executors.eks.utils import ( + CONFIG_DEFAULTS, + CONFIG_GROUP_NAME, + AllEksConfigKeys, +) +from airflow.providers.amazon.aws.hooks.eks import EksHook +from airflow.providers.amazon.aws.hooks.sts import StsHook +from airflow.providers.amazon.aws.utils.eks_get_token import ( + TOKEN_EXPIRATION_MINUTES, + fetch_access_token_for_cluster, +) +from airflow.providers.common.compat.sdk import conf + +if TYPE_CHECKING: + from kubernetes import client + from kubernetes_asyncio import client as async_client + +# UPDATING still serves the Kubernetes API, so every other status is refused. +_USABLE_CLUSTER_STATUSES = ("ACTIVE", "UPDATING") + + +def _get_sync_token_key(kubernetes_version: str) -> str: + # kubernetes 36 moved the bearer token from the "authorization" api_key to "BearerToken". + return "BearerToken" if Version(kubernetes_version).major >= 36 else "authorization" + + +_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.""" + from kubernetes import client + + configuration = client.Configuration() + _configure_eks_auth(configuration, _SYNC_TOKEN_KEY) + return client.CoreV1Api(client.ApiClient(configuration=configuration)) + + +def _get_eks_async_kube_client() -> 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") + return async_client.CoreV1Api(async_client.ApiClient(configuration=configuration)) + + +def _configure_eks_auth( + configuration: client.Configuration | async_client.Configuration, token_key: str +) -> None: + cluster_name = conf.get(CONFIG_GROUP_NAME, AllEksConfigKeys.CLUSTER_NAME, fallback=None) + 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." + ) + region_name = conf.get(CONFIG_GROUP_NAME, AllEksConfigKeys.REGION_NAME, fallback=None) + conn_id = conf.get( + CONFIG_GROUP_NAME, + AllEksConfigKeys.AWS_CONN_ID, + fallback=CONFIG_DEFAULTS[AllEksConfigKeys.AWS_CONN_ID], + ) + + eks_hook = EksHook(aws_conn_id=conn_id, region_name=region_name) + cluster = eks_hook.conn.describe_cluster(name=cluster_name)["cluster"] + if cluster["status"] not in _USABLE_CLUSTER_STATUSES: + raise ValueError( + f"EKS cluster {cluster_name} is {cluster['status']}; the executor needs it to be ACTIVE or " + f"UPDATING. Wait for it to become ACTIVE, or point [{CONFIG_GROUP_NAME}] cluster_name at a " + "usable cluster." + ) + session = eks_hook.get_session() + + # EKS only accepts tokens presigned against the regional STS endpoint. + os.environ["AWS_STS_REGIONAL_ENDPOINTS"] = "regional" + try: + sts_endpoint = StsHook( + aws_conn_id=conn_id, region_name=session.region_name + ).conn_client_meta.endpoint_url + finally: + del os.environ["AWS_STS_REGIONAL_ENDPOINTS"] Review Comment: I'm not sure if anything else uses "AWS_STS_REGIONAL_ENDPOINTS". Is it possible that this had a value before L108 that we should restore here instead of deleting? ########## providers/amazon/src/airflow/providers/amazon/aws/executors/eks/eks_executor.py: ########## @@ -0,0 +1,137 @@ +# 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. +"""AwsEksExecutor: run Airflow tasks as pods on an Amazon EKS cluster.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +from airflow.providers.amazon.aws.executors.eks.utils import ( + CONFIG_DEFAULTS, + CONFIG_GROUP_NAME, + AllEksConfigKeys, +) +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, conf + +_CNCF_KUBERNETES_REQUIRED = "AwsEksExecutor requires apache-airflow-providers-cncf-kubernetes>=10.24.0" + +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", {}) +): + 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" + + +class AwsEksExecutor(KubernetesExecutor): + """ + A KubernetesExecutor that authenticates against an Amazon EKS cluster. + + Builds the Kubernetes client from ``[aws_eks_executor]`` configuration and keeps the + short-lived EKS token fresh in-process. All pod-level behaviour comes unchanged from the + KubernetesExecutor and its ``[kubernetes_executor]`` configuration. + """ + + supports_multi_team: bool = True + + def __init__(self, *args, **kwargs): + self._validate_eks_config() + super().__init__(*args, **kwargs) + # Needs team_name, which super().__init__ sets. + 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( + CONFIG_GROUP_NAME, + AllEksConfigKeys.CHECK_HEALTH_ON_STARTUP, + fallback=CONFIG_DEFAULTS[AllEksConfigKeys.CHECK_HEALTH_ON_STARTUP], + ) + + if not check_health: + return + + self.log.info("Starting EKS Executor and determining health...") + try: + self.check_health() + except RuntimeError: + self.log.error("Stopping the Airflow Scheduler from starting until the issue is resolved.") + raise + + def check_health(self) -> None: + """ + Make a test Kubernetes API call to check the health of the EKS Executor. + + Building the client only proves the cluster exists; without this, a missing EKS access + entry or RBAC binding only shows up later as a watcher error loop. + """ + if TYPE_CHECKING: + assert self.kube_client + namespace = self.kube_config.kube_namespace + try: + self.kube_client.list_namespaced_pod(namespace, limit=1) + except ApiException as e: + raise RuntimeError( + f"EKS Executor health check has failed because: cannot list pods in namespace {namespace} " + f"({e.status} {e.reason}). The IAM principal behind the AWS connection needs an EKS access " + "entry (or aws-auth mapping) on the cluster that allows create, get, list, watch, patch and " + f"delete on pods and get on pods/log in namespace {namespace}. See 'Grant access to the " + "cluster' in the AwsEksExecutor docs." + ) 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): + 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 _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. + os.environ[f"{env_var_prefix}KUBERNETES_EXECUTOR__{key.upper()}"] = path Review Comment: Won't this cause conflicts in a multi-executor setup with EKS and K8S executors running side by side? It looks liken we're stealing their `KUBERNETES_EXECUTOR__` namespace and overwriting it. ########## providers/amazon/tests/unit/amazon/aws/executors/eks/test_eks_executor.py: ########## @@ -0,0 +1,154 @@ +# 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 + +import importlib +import os +import sys +from unittest import mock + +import pytest +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.executors.kubernetes_executor import KubernetesExecutor +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, conf + +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" +CLIENT_FACTORY_ENV_VAR = "AIRFLOW__KUBERNETES_EXECUTOR__CLIENT_FACTORY" +TEAMS = [ + pytest.param(None, "AIRFLOW__", id="global"), + pytest.param( + "team_a", + "AIRFLOW__TEAM_A___", + id="team", + marks=pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="Multi-team requires Airflow 3.2+"), + ), +] + + +def _team_kwargs(team_name: str | None) -> dict: + return {"team_name": team_name} if team_name else {} + + [email protected](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) + yield + + +class TestAwsEksExecutor: Review Comment: `test_team_config` in `test_ecs_executor.py` shows both global and team-scoped values getting tested and I don't see the equivalent here, can you add them? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
