diff --git a/providers/amazon/docs/operators/eks.rst b/providers/amazon/docs/operators/eks.rst index 76c0f5ea4ed11..1b94e31789533 100644 --- a/providers/amazon/docs/operators/eks.rst +++ b/providers/amazon/docs/operators/eks.rst @@ -205,6 +205,26 @@ Note: An Amazon EKS Cluster with underlying compute infrastructure is required. :start-after: [START howto_operator_eks_pod_operator] :end-before: [END howto_operator_eks_pod_operator] +.. _howto/decorator:eks_pod: + +Run a Python task on an Amazon EKS Cluster from TaskFlow +======================================================== + +The ``@task.eks_pod`` decorator runs a decorated Python function inside a pod on an existing Amazon +EKS Cluster, using the TaskFlow style. It wraps +:class:`~airflow.providers.amazon.aws.operators.eks.EksPodOperator`, which builds the cluster +kubeconfig and a short-lived token from the AWS connection (``aws_conn_id``) at run time. This means +the worker does not need the ``aws`` CLI or a pre-built kubeconfig, which is what running +``@task.kubernetes`` against EKS would otherwise require. + +An Amazon EKS Cluster with underlying compute infrastructure is required. + +.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_eks_pod_decorator.py + :language: python + :dedent: 4 + :start-after: [START howto_decorator_eks_pod] + :end-before: [END howto_decorator_eks_pod] + Sensors ------- diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml index 3d26682c4d5ec..ea1a1149a0174 100644 --- a/providers/amazon/provider.yaml +++ b/providers/amazon/provider.yaml @@ -986,6 +986,10 @@ transfers: how-to-guide: /docs/apache-airflow-providers-amazon/transfer/s3_to_dynamodb.rst python-module: airflow.providers.amazon.aws.transfers.s3_to_dynamodb +task-decorators: + - class-name: airflow.providers.amazon.aws.decorators.eks.eks_pod_task + name: eks_pod + extra-links: - airflow.providers.amazon.aws.links.athena.AthenaQueryResultsLink - airflow.providers.amazon.aws.links.batch.BatchJobDefinitionLink diff --git a/providers/amazon/src/airflow/providers/amazon/aws/decorators/__init__.py b/providers/amazon/src/airflow/providers/amazon/aws/decorators/__init__.py new file mode 100644 index 0000000000000..21d298ede6ed3 --- /dev/null +++ b/providers/amazon/src/airflow/providers/amazon/aws/decorators/__init__.py @@ -0,0 +1,17 @@ +# 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 diff --git a/providers/amazon/src/airflow/providers/amazon/aws/decorators/eks.py b/providers/amazon/src/airflow/providers/amazon/aws/decorators/eks.py new file mode 100644 index 0000000000000..af219944faf30 --- /dev/null +++ b/providers/amazon/src/airflow/providers/amazon/aws/decorators/eks.py @@ -0,0 +1,172 @@ +# 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 base64 +import os +import pickle +from collections.abc import Callable, Sequence +from shlex import quote +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING + +import dill +from kubernetes.client import models as k8s + +from airflow.providers.amazon.aws.operators.eks import EksPodOperator +from airflow.providers.cncf.kubernetes.python_kubernetes_script import write_python_script +from airflow.providers.common.compat.sdk import ( + DecoratedOperator, + TaskDecorator, + task_decorator_factory, +) + +if TYPE_CHECKING: + from airflow.sdk import Context + +_PYTHON_SCRIPT_ENV = "__PYTHON_SCRIPT" +_PYTHON_INPUT_ENV = "__PYTHON_INPUT" + + +def _generate_decoded_command(env_var: str, file: str) -> str: + return ( + f'python -c "import base64, os;' + rf"x = base64.b64decode(os.environ[\"{env_var}\"]);" + rf'f = open(\"{file}\", \"wb\"); f.write(x); f.close()"' + ) + + +def _read_file_contents(filename: str) -> str: + with open(filename, "rb") as script_file: + return base64.b64encode(script_file.read()).decode("utf-8") + + +class _EksPodDecoratedOperator(DecoratedOperator, EksPodOperator): + """Wrap a Python callable to run inside a pod on Amazon EKS via ``EksPodOperator``.""" + + custom_operator_name = "@task.eks_pod" + + # `cmds` and `arguments` are used internally by the operator + template_fields: Sequence[str] = tuple( + {"op_args", "op_kwargs", *EksPodOperator.template_fields} - {"cmds", "arguments"} + ) + + # Since we won't mutate the arguments, we should just do the shallow copy + # there are some cases we can't deepcopy the objects (e.g protobuf). + shallow_copy_attrs: Sequence[str] = ("python_callable",) + + def __init__(self, *, cluster_name: str, use_dill: bool = False, **kwargs) -> None: + self.use_dill = use_dill + + # Accept the EKS ``pod_name`` or the K8s-style ``name``, otherwise derive one from the callable. + pod_name = ( + kwargs.pop("pod_name", None) + or kwargs.pop("name", None) + or f"eks-airflow-pod-{kwargs['python_callable'].__name__}" + ) + random_name_suffix = kwargs.pop("random_name_suffix", True) + super().__init__( + cluster_name=cluster_name, + pod_name=pod_name, + random_name_suffix=random_name_suffix, + cmds=["placeholder-command"], + **kwargs, + ) + + def _generate_cmds(self) -> list[str]: + script_filename = "/tmp/script.py" + input_filename = "/tmp/script.in" + + if getattr(self, "do_xcom_push", False): + output_filename = "/airflow/xcom/return.json" + make_xcom_dir_cmd = "mkdir -p /airflow/xcom" + else: + output_filename = "/dev/null" + make_xcom_dir_cmd = ":" # shell no-op + + write_local_script_file_cmd = ( + f"{_generate_decoded_command(quote(_PYTHON_SCRIPT_ENV), quote(script_filename))}" + ) + write_local_input_file_cmd = ( + f"{_generate_decoded_command(quote(_PYTHON_INPUT_ENV), quote(input_filename))}" + ) + exec_python_cmd = f"python {script_filename} {input_filename} {output_filename}" + return [ + "bash", + "-cx", + ( + f"{write_local_script_file_cmd} && " + f"{write_local_input_file_cmd} && " + f"{make_xcom_dir_cmd} && " + f"{exec_python_cmd}" + ), + ] + + def execute(self, context: Context): + with TemporaryDirectory(prefix="venv") as tmp_dir: + pickling_library = dill if self.use_dill else pickle + script_filename = os.path.join(tmp_dir, "script.py") + input_filename = os.path.join(tmp_dir, "script.in") + + with open(input_filename, "wb") as file: + pickling_library.dump({"args": self.op_args, "kwargs": self.op_kwargs}, file) + + py_source = self.get_python_source() + jinja_context = { + "op_args": self.op_args, + "op_kwargs": self.op_kwargs, + "pickling_library": pickling_library.__name__, + "python_callable": self.python_callable.__name__, + "python_callable_source": py_source, + "string_args_global": False, + } + write_python_script(jinja_context=jinja_context, filename=script_filename) + + self.env_vars: list[k8s.V1EnvVar] = [ + *self.env_vars, + k8s.V1EnvVar(name=_PYTHON_SCRIPT_ENV, value=_read_file_contents(script_filename)), + k8s.V1EnvVar(name=_PYTHON_INPUT_ENV, value=_read_file_contents(input_filename)), + ] + + self.cmds = self._generate_cmds() + return super().execute(context) + + +def eks_pod_task( + python_callable: Callable | None = None, + multiple_outputs: bool | None = None, + **kwargs, +) -> TaskDecorator: + """ + Run a Python function in a pod on an Amazon EKS cluster. + + Wraps :class:`~airflow.providers.amazon.aws.operators.eks.EksPodOperator`, which builds the + cluster kubeconfig and a short-lived token from the AWS connection (``aws_conn_id``), so the + worker needs no ``aws`` CLI or kubeconfig. Any ``EksPodOperator`` argument is accepted via + ``kwargs``. + + :param python_callable: Function to decorate + :param multiple_outputs: if set, function return value will be + unrolled to multiple XCom values. Dict will unroll to xcom values with + keys as XCom keys. Defaults to False. + """ + return task_decorator_factory( + python_callable=python_callable, + multiple_outputs=multiple_outputs, + decorated_operator_class=_EksPodDecoratedOperator, + **kwargs, + ) diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py index ea3e1f49437cc..461d2ff3e1d75 100644 --- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py +++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py @@ -1155,6 +1155,9 @@ def get_provider_info(): "python-module": "airflow.providers.amazon.aws.transfers.s3_to_dynamodb", }, ], + "task-decorators": [ + {"class-name": "airflow.providers.amazon.aws.decorators.eks.eks_pod_task", "name": "eks_pod"} + ], "extra-links": [ "airflow.providers.amazon.aws.links.athena.AthenaQueryResultsLink", "airflow.providers.amazon.aws.links.batch.BatchJobDefinitionLink", diff --git a/providers/amazon/tests/system/amazon/aws/example_eks_pod_decorator.py b/providers/amazon/tests/system/amazon/aws/example_eks_pod_decorator.py new file mode 100644 index 0000000000000..4b7f35f23a548 --- /dev/null +++ b/providers/amazon/tests/system/amazon/aws/example_eks_pod_decorator.py @@ -0,0 +1,213 @@ +# +# 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 datetime import datetime + +import boto3 + +from airflow.providers.amazon.aws.hooks.eks import ClusterStates, NodegroupStates +from airflow.providers.amazon.aws.operators.eks import ( + EksCreateClusterOperator, + EksCreateNodegroupOperator, + EksDeleteClusterOperator, + EksDeleteNodegroupOperator, +) +from airflow.providers.amazon.aws.sensors.eks import EksClusterStateSensor, EksNodegroupStateSensor + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import DAG, chain, task +else: + # Airflow 2 path + from airflow.decorators import task # type: ignore[attr-defined,no-redef] + from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef] + from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment] + +try: + from airflow.sdk import TriggerRule +except ImportError: + # Compatibility for Airflow < 3.1 + from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] + +from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +from system.amazon.aws.utils.k8s import get_describe_pod_operator + +DAG_ID = "example_eks_pod_decorator" + +# Externally fetched variables: +ROLE_ARN_KEY = "ROLE_ARN" +SUBNETS_KEY = "SUBNETS" + +sys_test_context_task = ( + SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).add_variable(SUBNETS_KEY, split_string=True).build() +) + + +@task +def create_launch_template(template_name: str): + # This launch template enables IMDSv2. + boto3.client("ec2").create_launch_template( + LaunchTemplateName=template_name, + LaunchTemplateData={ + "MetadataOptions": {"HttpEndpoint": "enabled", "HttpTokens": "required"}, + }, + ) + + +@task(trigger_rule=TriggerRule.ALL_DONE) +def delete_launch_template(template_name: str): + boto3.client("ec2").delete_launch_template(LaunchTemplateName=template_name) + + +with DAG( + dag_id=DAG_ID, + schedule="@once", + start_date=datetime(2021, 1, 1), + catchup=False, +) as dag: + test_context = sys_test_context_task() + env_id = test_context[ENV_ID_KEY] + + cluster_name = f"{env_id}-cluster" + nodegroup_name = f"{env_id}-nodegroup" + launch_template_name = f"{env_id}-launch-template" + + create_cluster = EksCreateClusterOperator( + task_id="create_cluster", + cluster_name=cluster_name, + cluster_role_arn=test_context[ROLE_ARN_KEY], + resources_vpc_config={"subnetIds": test_context[SUBNETS_KEY]}, + compute=None, + ) + + await_create_cluster = EksClusterStateSensor( + task_id="await_create_cluster", + cluster_name=cluster_name, + target_state=ClusterStates.ACTIVE, + ) + + create_nodegroup = EksCreateNodegroupOperator( + task_id="create_nodegroup", + cluster_name=cluster_name, + nodegroup_name=nodegroup_name, + nodegroup_subnets=test_context[SUBNETS_KEY], + nodegroup_role_arn=test_context[ROLE_ARN_KEY], + ) + # The launch template enforces IMDSv2 and is required for internal compliance when running + # these system tests on AWS infrastructure. It is not required for the decorator to work. + create_nodegroup.create_nodegroup_kwargs = {"launchTemplate": {"name": launch_template_name}} + + await_create_nodegroup = EksNodegroupStateSensor( + task_id="await_create_nodegroup", + cluster_name=cluster_name, + nodegroup_name=nodegroup_name, + target_state=NodegroupStates.ACTIVE, + poke_interval=10, + ) + + # [START howto_decorator_eks_pod] + # Credentials come from the AWS connection; the worker needs no aws CLI or kubeconfig. + @task.eks_pod( + pod_name="run-pod", + cluster_name=cluster_name, + image="python:3.12-slim", + get_logs=True, + on_finish_action="keep_pod", + ) + def run_pod(): + print("Hello from an EKS pod") + + start_pod = run_pod() + # [END howto_decorator_eks_pod] + + describe_pod = get_describe_pod_operator( + cluster_name, pod_name="{{ ti.xcom_pull(key='pod_name', task_ids='run_pod') }}" + ) + # only describe the pod if the task above failed, to help diagnose + describe_pod.trigger_rule = TriggerRule.ONE_FAILED + + await_nodegroup_stable = EksNodegroupStateSensor( + task_id="await_nodegroup_stable", + trigger_rule=TriggerRule.ALL_DONE, + cluster_name=cluster_name, + nodegroup_name=nodegroup_name, + target_state=NodegroupStates.ACTIVE, + ) + + delete_nodegroup = EksDeleteNodegroupOperator( + task_id="delete_nodegroup", + cluster_name=cluster_name, + nodegroup_name=nodegroup_name, + trigger_rule=TriggerRule.ALL_DONE, + ) + + await_delete_nodegroup = EksNodegroupStateSensor( + task_id="await_delete_nodegroup", + trigger_rule=TriggerRule.ALL_DONE, + cluster_name=cluster_name, + nodegroup_name=nodegroup_name, + target_state=NodegroupStates.NONEXISTENT, + ) + + delete_cluster = EksDeleteClusterOperator( + task_id="delete_cluster", + cluster_name=cluster_name, + trigger_rule=TriggerRule.ALL_DONE, + ) + + await_delete_cluster = EksClusterStateSensor( + task_id="await_delete_cluster", + trigger_rule=TriggerRule.ALL_DONE, + cluster_name=cluster_name, + target_state=ClusterStates.NONEXISTENT, + poke_interval=10, + ) + + chain( + # TEST SETUP + test_context, + create_launch_template(launch_template_name), + # TEST BODY + create_cluster, + await_create_cluster, + create_nodegroup, + await_create_nodegroup, + start_pod, + # TEST TEARDOWN + describe_pod, + await_nodegroup_stable, + delete_nodegroup, + await_delete_nodegroup, + delete_cluster, + await_delete_cluster, + delete_launch_template(launch_template_name), + ) + + from tests_common.test_utils.watcher import watcher + + # This test needs watcher in order to properly mark success/failure + # when "tearDown" task with trigger rule is part of the DAG + list(dag.tasks) >> watcher() + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/amazon/tests/unit/amazon/aws/decorators/__init__.py b/providers/amazon/tests/unit/amazon/aws/decorators/__init__.py new file mode 100644 index 0000000000000..21d298ede6ed3 --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/decorators/__init__.py @@ -0,0 +1,17 @@ +# 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 diff --git a/providers/amazon/tests/unit/amazon/aws/decorators/test_eks.py b/providers/amazon/tests/unit/amazon/aws/decorators/test_eks.py new file mode 100644 index 0000000000000..f18ca4c37241e --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/decorators/test_eks.py @@ -0,0 +1,97 @@ +# 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 tests_common.test_utils.compat import timezone +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import DAG +else: + from airflow.models.dag import DAG # type: ignore[no-redef,assignment] + +from airflow.providers.amazon.aws.decorators.eks import _EksPodDecoratedOperator, eks_pod_task + +DEFAULT_DATE = timezone.datetime(2023, 1, 1) + + +class TestEksPodDecorator: + def test_init_builds_eks_pod_operator(self): + with DAG(dag_id="test_eks_deco_init", schedule=None, start_date=DEFAULT_DATE): + + @eks_pod_task(cluster_name="my-eks", image="python:3.12-slim") + def f(): + return {"a": 1} + + task = f() + + op = task.operator + assert isinstance(op, _EksPodDecoratedOperator) + assert op.custom_operator_name == "@task.eks_pod" + assert op.task_id == "f" + assert op.cluster_name == "my-eks" + assert op.image == "python:3.12-slim" + assert op.cmds == ["placeholder-command"] + assert op.random_name_suffix is True + assert op.pod_name == "eks-airflow-pod-f" + + def test_explicit_pod_name_is_respected(self): + with DAG(dag_id="test_eks_deco_name", schedule=None, start_date=DEFAULT_DATE): + + @eks_pod_task(cluster_name="my-eks", image="python:3.12-slim", pod_name="custom-pod") + def f(): + return None + + task = f() + + assert task.operator.pod_name == "custom-pod" + + @mock.patch("airflow.providers.amazon.aws.decorators.eks.EksPodOperator.execute") + def test_execute_injects_script_then_delegates(self, mock_super_execute): + with DAG(dag_id="test_eks_deco_exec", schedule=None, start_date=DEFAULT_DATE): + + @eks_pod_task(cluster_name="my-eks", image="python:3.12-slim") + def f(): + return {"a": 1} + + task = f() + + op = task.operator + op.execute(context=mock.MagicMock()) + + env_names = {env.name for env in op.env_vars} + assert {"__PYTHON_SCRIPT", "__PYTHON_INPUT"} <= env_names + assert op.cmds[0] == "bash" + mock_super_execute.assert_called_once() + + @pytest.mark.parametrize("do_xcom_push", [True, False]) + def test_generate_cmds_handles_xcom(self, do_xcom_push): + with DAG(dag_id="test_eks_deco_xcom", schedule=None, start_date=DEFAULT_DATE): + + @eks_pod_task(cluster_name="my-eks", image="python:3.12-slim", do_xcom_push=do_xcom_push) + def f(): + return {"a": 1} + + task = f() + + cmds = task.operator._generate_cmds() + assert cmds[:2] == ["bash", "-cx"] + assert ("mkdir -p /airflow/xcom" in cmds[2]) is do_xcom_push