seanghaeli commented on code in PR #74399: URL: https://github.com/apache/airflow/pull/74399#discussion_r4291655987
########## 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: It's a good point to bring up, and it's something I did consider explicitly when writing #73014. To keep the changes to `cncf` minimal, I decided to bring `AwsEksExecutor` to multi-team feature parity with `KubernetesExecutor` ([docs](https://github.com/apache/airflow/blob/313255f8db6543f8482bd1c391d97a46567da89f/providers/amazon/docs/executors/eks-executor.rst?plain=1#L203-L207)), which reads its cluster settings globally, so teams share one cluster ([kube_client.py](https://github.com/apache/airflow/blob/313255f8db6543f8482bd1c391d97a46567da89f/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_client.py#L180-L181)): ~~~python if in_cluster is None: in_cluster = conf.getboolean("kubernetes_executor", "in_cluster") ~~~ On the other hand, the AWS executors read their whole section per team, so each team can have its own cluster ([ecs_executor.py](https://github.com/apache/airflow/blob/313255f8db6543f8482bd1c391d97a46567da89f/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py#L142)): ~~~python self.cluster = self.conf.get(CONFIG_GROUP_NAME, AllEcsConfigKeys.CLUSTER) ~~~ Getting EKS to that level needs cncf to pass the team name to the client factory, which is a small change I'll do as a follow-up. I'll file an issue for it and target for it to be included in the next cncf release. If you, @o-nikolas, and @vincbeck are happy with that plan, then we'll treat this as a follow-on. -- 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]
