Kubernetes decorator

airflow/decorators/__init__.pyi

@@ -188,8 +188,7 @@ class TaskDecoratorCollection:
             Dict will unroll to XCom values with keys as XCom keys. Defaults to False.
         :param use_dill: Whether to use dill or pickle for serialization
         :param python_command: Python command for executing functions, Default: python3
-        :param image: Docker image from which to create the container.
-            If image tag is omitted, "latest" will be used.
+        :param image: Docker image from which to create the container.(Required)
         :param api_version: Remote API version. Set to ``auto`` to automatically
             detect the server's version.
         :param container_name: Name of the container. Optional (templated)
@@ -245,5 +244,119 @@ class TaskDecoratorCollection:
         :param cap_add: Include container capabilities
         """
         # [END decorator_signature]
+    def kubernetes(
+        self,
+        *,
+        multiple_outputs: Optional[bool] = None,
+        namespace: Optional[str] = None,
+        image: str,
+        name: Optional[str] = None,
+        random_name_suffix: bool = True,
+        cmds: Optional[List[str]] = None,
+        arguments: Optional[List[str]] = None,
+        ports: Optional[List[k8s.V1ContainerPort]] = None,
+        volume_mounts: Optional[List[k8s.V1VolumeMount]] = None,
+        volumes: Optional[List[k8s.V1Volume]] = None,
+        env_vars: Optional[List[k8s.V1EnvVar]] = None,
+        env_from: Optional[List[k8s.V1EnvFromSource]] = None,
+        secrets: Optional[List[Secret]] = None,
+        in_cluster: Optional[bool] = None,
+        cluster_context: Optional[str] = None,
+        labels: Optional[Dict] = None,
+        reattach_on_restart: bool = True,
+        startup_timeout_seconds: int = 120,
+        get_logs: bool = True,
+        image_pull_policy: Optional[str] = None,
+        annotations: Optional[Dict] = None,
+        resources: Optional[k8s.V1ResourceRequirements] = None,
+        affinity: Optional[k8s.V1Affinity] = None,
+        config_file: Optional[str] = None,
+        node_selectors: Optional[dict] = None,
+        node_selector: Optional[dict] = None,
+        image_pull_secrets: Optional[List[k8s.V1LocalObjectReference]] = None,
+        service_account_name: Optional[str] = None,
+        is_delete_operator_pod: bool = True,
+        hostnetwork: bool = False,
+        tolerations: Optional[List[k8s.V1Toleration]] = None,
+        security_context: Optional[Dict] = None,
+        dnspolicy: Optional[str] = None,
+        schedulername: Optional[str] = None,
+        full_pod_spec: Optional[k8s.V1Pod] = None,
+        init_containers: Optional[List[k8s.V1Container]] = None,
+        log_events_on_failure: bool = False,
+        do_xcom_push: bool = False,
+        pod_template_file: Optional[str] = None,
+        priority_class_name: Optional[str] = None,
+        pod_runtime_info_envs: Optional[List[PodRuntimeInfoEnv]] = None,
+        termination_grace_period: Optional[int] = None,
+        configmaps: Optional[List[str]] = None,
+        **kwargs,
+    ) -> TaskDecorator:
+        """Wraps a function to be executed on a k8s pod using KubernetesPodOperator
+
+        Also accepts any argument that KubernetesPodOperator will via ``kwargs``.
+
+        :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.
+        :param namespace: the namespace to run within kubernetes.
+        :param image: Docker image you wish to launch. Defaults to hub.docker.com,
+            but fully qualified URLS will point to custom repositories. (templated)
+        :param name: name of the pod in which the task will run, will be used (plus a random
+            suffix if random_name_suffix is True) to generate a pod id (DNS-1123 subdomain,
+            containing only [a-z0-9.-]).
+        :param random_name_suffix: if True, will generate a random suffix.
+        :param cmds: entrypoint of the container. (templated)
+            The docker images's entrypoint is used if this is not provided.
+        :param arguments: arguments of the entrypoint. (templated)
+            The docker image's CMD is used if this is not provided.
+        :param ports: ports for launched pod.
+        :param volume_mounts: volumeMounts for launched pod.
+        :param volumes: volumes for launched pod. Includes ConfigMaps and PersistentVolumes.
+        :param env_vars: Environment variables initialized in the container. (templated)
+        :param secrets: Kubernetes secrets to inject in the container.
+            They can be exposed as environment vars or files in a volume.
+        :param in_cluster: run kubernetes client with in_cluster configuration.
+        :param cluster_context: context that points to kubernetes cluster.
+            Ignored when in_cluster is True. If None, current-context is used.
+        :param reattach_on_restart: if the scheduler dies while the pod is running, reattach and monitor
+        :param labels: labels to apply to the Pod. (templated)
+        :param startup_timeout_seconds: timeout in seconds to startup the pod.
+        :param get_logs: get the stdout of the container as logs of the tasks.
+        :param image_pull_policy: Specify a policy to cache or always pull an image.
+        :param annotations: non-identifying metadata you can attach to the Pod.
+            Can be a large range of data, and can include characters
+            that are not permitted by labels.
+        :param resources: A dict containing resources requests and limits.
+            Possible keys are request_memory, request_cpu, limit_memory, limit_cpu,
+            and limit_gpu, which will be used to generate airflow.kubernetes.pod.Resources.
+            See also kubernetes.io/docs/concepts/configuration/manage-compute-resources-container
+        :param affinity: A dict containing a group of affinity scheduling rules.
+        :param config_file: The path to the Kubernetes config file. (templated)
+            If not specified, default value is ``~/.kube/config``
+        :param node_selector: A dict containing a group of scheduling rules.
+        :param image_pull_secrets: Any image pull secrets to be given to the pod.
+            If more than one secret is required, provide a
+            comma separated list: secret_a,secret_b
+        :param service_account_name: Name of the service account
+        :param is_delete_operator_pod: What to do when the pod reaches its final
+            state, or the execution is interrupted. If True (default), delete the
+            pod; if False, leave the pod.
+        :param hostnetwork: If True enable host networking on the pod.
+        :param tolerations: A list of kubernetes tolerations.
+        :param security_context: security options the pod should run with (PodSecurityContext).
+        :param dnspolicy: dnspolicy for the pod.
+        :param schedulername: Specify a schedulername for the pod
+        :param full_pod_spec: The complete podSpec
+        :param init_containers: init container for the launched Pod
+        :param log_events_on_failure: Log the pod's events if a failure occurs
+        :param do_xcom_push: If True, the content of the file
+            /airflow/xcom/return.json in the container will also be pushed to an
+            XCom when the container completes.
+        :param pod_template_file: path to pod template file (templated)
+        :param priority_class_name: priority class name for the launched Pod
+        :param termination_grace_period: Termination grace period if task killed in UI,
+            defaults to kubernetes default
+        """
 
 task: TaskDecoratorCollection

airflow/example_dags/example_kubernetes_decorator.py

@@ -0,0 +1,57 @@
+#
+# 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.
+"""
+This is an example dag for using the KubernetesPodOperator.
+"""
+
+from datetime import datetime
+
+from airflow import DAG
+from airflow.decorators import task
+
+with DAG(
+    dag_id='example_kubernetes_decorator',
+    schedule_interval=None,
+    start_date=datetime(2021, 1, 1),
+    tags=['example'],
+    catchup=False,
+) as dag:
+
+    @task.kubernetes(image='python:3.8-slim-buster', name='k8s_test', namespace='default')
+    def execute_in_k8s_pod():
+        import time
+
+        print("Hello from k8s pod")
+        time.sleep(2)
+
+    @task.kubernetes(image='python:3.8-slim-buster')
+    def print_pattern():
+        n = 5
+        for i in range(0, n):
+            # inner loop to handle number of columns
+            # values changing acc. to outer loop
+            for j in range(0, i + 1):
+                # printing stars
+                print("* ", end="")
+
+            # ending line after each row
+            print("\r")
+
+    execute_in_k8s_pod_instance = execute_in_k8s_pod()
+    print_pattern_instance = print_pattern()
+    execute_in_k8s_pod_instance >> print_pattern_instance

airflow/providers/cncf/kubernetes/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.

airflow/providers/cncf/kubernetes/decorators/kubernetes.py

@@ -0,0 +1,145 @@
+# 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.
+
+import inspect
+import os
+import pickle
+import uuid
+from tempfile import TemporaryDirectory
+from textwrap import dedent
+from typing import TYPE_CHECKING, Callable, Optional, Sequence, TypeVar
+
+from kubernetes.client import models as k8s
+
+from airflow.decorators.base import DecoratedOperator, TaskDecorator, task_decorator_factory
+from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
+from airflow.utils.python_virtualenv import remove_task_decorator, write_python_script
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+def _generate_decode_command(env_var, file):
+    # We don't need `f.close()` as the interpreter is about to exit anyway
+    return (
+        f'python -c "import base64, os;'
+        rf'x = os.environ[\"{env_var}\"];'
+        rf'f = open(\"{file}\", \"w\"); f.write(x);"'
+    )
+
+
+def _read_file_contents(filename):
+    with open(filename) as script_file:
+        return script_file.read()
+
+
+class _KubernetesDecoratedOperator(DecoratedOperator, KubernetesPodOperator):
+    """
+    Wraps a Python callable and executes in a kubernetes pod
+
+    :param python_callable: A reference to an object that is callable
+    :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+        in your function (templated)
+    :param op_args: a list of positional arguments that will get unpacked when
+        calling your callable (templated)
+    :param multiple_outputs: if set, function return value will be
+        unrolled to multiple XCom values. Dict will unroll to xcom values with keys as keys.
+        Defaults to False.
+    """
+
+    template_fields: Sequence[str] = ('op_args', 'op_kwargs')
+
+    # 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,
+        **kwargs,
+    ) -> None:
+        self.pickling_library = pickle
+
+        # Set defaults for name and namespace.
+        if 'name' not in kwargs:
+            kwargs['name'] = f'k8s_airflow_pod_{uuid.uuid4().hex}'
+
+        if 'namespace' not in kwargs:
+            kwargs['namespace'] = 'default'
+
+        super().__init__(**kwargs)
+
+    def execute(self, context: 'Context'):
+
+        with TemporaryDirectory(prefix='venv') as tmp_dir:
+            script_filename = os.path.join(tmp_dir, 'script.py')
+            py_source = self._get_python_source()
+
+            jinja_context = dict(
+                op_args=self.op_args,
+                op_kwargs=self.op_kwargs,
+                pickling_library=self.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,
+                template_file='python_kubernetes_script.jinja2',
+            )
+
+            self.env_vars.append(
+                k8s.V1EnvVar(name="__PYTHON_SCRIPT", value=_read_file_contents(script_filename))
+            )
+
+            self.cmds.append("bash")
+
+            self.arguments.append("-cx")
+            self.arguments.append(
+                f'{_generate_decode_command("__PYTHON_SCRIPT", "/tmp/script.py")} && python /tmp/script.py'
+            )
+
+            return super().execute(context)
+
+    def _get_python_source(self):
+        raw_source = inspect.getsource(self.python_callable)
+        res = dedent(raw_source)
+        res = remove_task_decorator(res, "@task.kubernetes")
+        return res
+
+
+T = TypeVar("T", bound=Callable)
+
+
+def kubernetes_task(
+    python_callable: Optional[Callable] = None, multiple_outputs: Optional[bool] = None, **kwargs
+) -> TaskDecorator:
+    """
+    Kubernetes operator decorator. Wraps a function to be executed in K8s using KubernetesPodOperator.
+    Also accepts any argument that DockerOperator will via ``kwargs``. Can be reused in a single DAG.
+
+    :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=_KubernetesDecoratedOperator,
+        **kwargs,
+    )

airflow/providers/cncf/kubernetes/provider.yaml

@@ -78,3 +78,7 @@ hook-class-names:  # deprecated - to be removed after providers add dependency o
 connection-types:
   - hook-class-name: airflow.providers.cncf.kubernetes.hooks.kubernetes.KubernetesHook
     connection-type: kubernetes
+
+task-decorators:
+  - class-name: airflow.providers.cncf.kubernetes.decorators.kubernetes.kubernetes_task
+    name: kubernetes