You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by di...@apache.org on 2020/09/16 17:14:52 UTC
[airflow] branch v1-10-test updated: Add on_kill support for the
KubernetesPodOperator (#10666)
This is an automated email from the ASF dual-hosted git repository.
dimberman pushed a commit to branch v1-10-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v1-10-test by this push:
new efd0c3e Add on_kill support for the KubernetesPodOperator (#10666)
efd0c3e is described below
commit efd0c3ec6f838be4e5df022aa62bb0660503b367
Author: Daniel Imberman <da...@gmail.com>
AuthorDate: Wed Sep 2 07:48:29 2020 -0700
Add on_kill support for the KubernetesPodOperator (#10666)
This PR ensures that when a user kills a KubernetesPodOperator task
in the airflow UI, that the associated pod is also killed using the
on_kill method.
---
.../contrib/operators/kubernetes_pod_operator.py | 10 +
.../cncf/kubernetes/operators/kubernetes_pod.py} | 249 +++++++++++----------
kubernetes_tests/test_kubernetes_pod_operator.py | 28 +++
3 files changed, 166 insertions(+), 121 deletions(-)
diff --git a/airflow/contrib/operators/kubernetes_pod_operator.py b/airflow/contrib/operators/kubernetes_pod_operator.py
index 6ba57aa..fb984c0 100644
--- a/airflow/contrib/operators/kubernetes_pod_operator.py
+++ b/airflow/contrib/operators/kubernetes_pod_operator.py
@@ -229,6 +229,7 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
self.pod_template_file = pod_template_file
self.priority_class_name = priority_class_name
self.name = self._set_name(name)
+ self.client = None
@staticmethod
def create_labels_for_pod(context):
@@ -264,6 +265,8 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
client = kube_client.get_kube_client(cluster_context=self.cluster_context,
config_file=self.config_file)
+ self.client = client
+
# Add combination of labels to uniquely identify a running pod
labels = self.create_labels_for_pod(context)
@@ -444,3 +447,10 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
'Pod returned a failure: {state}'.format(state=final_state)
)
return final_state, result
+
+ def on_kill(self) -> None:
+ if self.pod:
+ pod = self.pod
+ namespace = pod.metadata.namespace
+ name = pod.metadata.name
+ self.client.delete_namespaced_pod(name=name, namespace=namespace, grace_period_seconds=0)
diff --git a/airflow/contrib/operators/kubernetes_pod_operator.py b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
similarity index 74%
copy from airflow/contrib/operators/kubernetes_pod_operator.py
copy to airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
index 6ba57aa..2785109 100644
--- a/airflow/contrib/operators/kubernetes_pod_operator.py
+++ b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
@@ -15,14 +15,20 @@
# specific language governing permissions and limitations
# under the License.
"""Executes task in a Kubernetes POD"""
-
import re
+from typing import Dict, Iterable, List, Optional, Tuple
+
import yaml
+from kubernetes.client import models as k8s
from airflow.exceptions import AirflowException
from airflow.kubernetes import kube_client, pod_generator, pod_launcher
from airflow.kubernetes.k8s_model import append_to_pod
-from airflow.kubernetes.pod import Resources
+from airflow.kubernetes.pod import Port, Resources
+from airflow.kubernetes.pod_runtime_info_env import PodRuntimeInfoEnv
+from airflow.kubernetes.secret import Secret
+from airflow.kubernetes.volume import Volume
+from airflow.kubernetes.volume_mount import VolumeMount
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.utils.helpers import validate_key
@@ -34,11 +40,18 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
"""
Execute a task in a Kubernetes Pod
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:KubernetesPodOperator`
+
.. note::
- If you use `Google Kubernetes Engine <https://cloud.google.com/kubernetes-engine/>`__, use
- :class:`~airflow.gcp.operators.kubernetes_engine.GKEPodOperator`, which
+ If you use `Google Kubernetes Engine <https://cloud.google.com/kubernetes-engine/>`__
+ and Airflow is not running in the same cluster, consider using
+ :class:`~airflow.providers.google.cloud.operators.kubernetes_engine.GKEStartPodOperator`, which
simplifies the authorization process.
+ :param namespace: the namespace to run within kubernetes.
+ :type namespace: str
:param image: Docker image you wish to launch. Defaults to hub.docker.com,
but fully qualified URLS will point to custom repositories.
:type image: str
@@ -51,25 +64,12 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
:param arguments: arguments of the entrypoint. (templated)
The docker image's CMD is used if this is not provided.
:type arguments: list[str]
- :param image_pull_policy: Specify a policy to cache or always pull an image.
- :type image_pull_policy: str
- :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
- :type image_pull_secrets: str
:param ports: ports for launched pod.
:type ports: list[airflow.kubernetes.pod.Port]
:param volume_mounts: volumeMounts for launched pod.
:type volume_mounts: list[airflow.kubernetes.volume_mount.VolumeMount]
:param volumes: volumes for launched pod. Includes ConfigMaps and PersistentVolumes.
:type volumes: list[airflow.kubernetes.volume.Volume]
- :param labels: labels to apply to the Pod.
- :type labels: dict
- :param startup_timeout_seconds: timeout in seconds to startup the pod.
- :type startup_timeout_seconds: int
- :param name: name of the pod in which the task will run, will be used to
- generate a pod id (DNS-1123 subdomain, containing only [a-z0-9.-]).
- :type name: str
:param env_vars: Environment variables initialized in the container. (templated)
:type env_vars: dict
:param secrets: Kubernetes secrets to inject in the container.
@@ -88,9 +88,11 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
:type startup_timeout_seconds: int
:param get_logs: get the stdout of the container as logs of the tasks.
:type get_logs: bool
+ :param image_pull_policy: Specify a policy to cache or always pull an image.
+ :type image_pull_policy: str
: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.
+ Can be a large range of data, and can include characters
+ that are not permitted by labels.
:type annotations: dict
:param resources: A dict containing resources requests and limits.
Possible keys are request_memory, request_cpu, limit_memory, limit_cpu,
@@ -99,16 +101,17 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
:type resources: dict
:param affinity: A dict containing a group of affinity scheduling rules.
:type affinity: dict
- :param node_selectors: A dict containing a group of scheduling rules.
- :type node_selectors: dict
- :param config_file: The path to the Kubernetes config file. (templated)
:param config_file: The path to the Kubernetes config file. (templated)
If not specified, default value is ``~/.kube/config``
:type config_file: str
- :param do_xcom_push: If do_xcom_push is True, the content of the file
- /airflow/xcom/return.json in the container will also be pushed to an
- XCom when the container completes.
- :type do_xcom_push: bool
+ :param node_selectors: A dict containing a group of scheduling rules.
+ :type node_selectors: dict
+ :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
+ :type image_pull_secrets: str
+ :param service_account_name: Name of the service account
+ :type service_account_name: str
:param is_delete_operator_pod: What to do when the pod reaches its final
state, or the execution is interrupted.
If False (default): do nothing, If True: delete the pod
@@ -120,11 +123,11 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
:param configmaps: A list of configmap names objects that we
want mount as env variables.
:type configmaps: list[str]
- :param pod_runtime_info_envs: environment variables about
- pod runtime information (ip, namespace, nodeName, podName).
- :type pod_runtime_info_envs: list[airflow.kubernetes.pod_runtime_info_env.PodRuntimeInfoEnv]
:param security_context: security options the pod should run with (PodSecurityContext).
:type security_context: dict
+ :param pod_runtime_info_envs: environment variables about
+ pod runtime information (ip, namespace, nodeName, podName).
+ :type pod_runtime_info_envs: list[airflow.kubernetes.pod_runtime_info_env.PodRuntimeInfoEnv]
:param dnspolicy: dnspolicy for the pod.
:type dnspolicy: str
:param schedulername: Specify a schedulername for the pod
@@ -141,54 +144,57 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
:type do_xcom_push: bool
:param pod_template_file: path to pod template file
:type pod_template_file: str
+ :param priority_class_name: priority class name for the launched Pod
+ :type priority_class_name: str
"""
- template_fields = ('image', 'cmds', 'arguments', 'env_vars', 'config_file', 'pod_template_file')
+ template_fields: Iterable[str] = (
+ 'image', 'cmds', 'arguments', 'env_vars', 'config_file', 'pod_template_file')
@apply_defaults
def __init__(self, # pylint: disable=too-many-arguments,too-many-locals
- namespace=None,
- image=None,
- name=None,
- cmds=None,
- arguments=None,
- ports=None,
- volume_mounts=None,
- volumes=None,
- env_vars=None,
- secrets=None,
- in_cluster=None,
- cluster_context=None,
- labels=None,
- reattach_on_restart=True,
- startup_timeout_seconds=120,
- get_logs=True,
- image_pull_policy='IfNotPresent',
- annotations=None,
- resources=None,
- affinity=None,
- config_file=None,
- node_selectors=None,
- image_pull_secrets=None,
- service_account_name='default',
- is_delete_operator_pod=False,
- hostnetwork=False,
- tolerations=None,
- configmaps=None,
- security_context=None,
- pod_runtime_info_envs=None,
- dnspolicy=None,
- schedulername=None,
- full_pod_spec=None,
- init_containers=None,
- log_events_on_failure=False,
- do_xcom_push=False,
- pod_template_file=None,
- priority_class_name=None,
- *args,
+ *,
+ namespace: Optional[str] = None,
+ image: Optional[str] = None,
+ name: Optional[str] = None,
+ cmds: Optional[List[str]] = None,
+ arguments: Optional[List[str]] = None,
+ ports: Optional[List[Port]] = None,
+ volume_mounts: Optional[List[VolumeMount]] = None,
+ volumes: Optional[List[Volume]] = None,
+ env_vars: Optional[Dict] = 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: str = 'IfNotPresent',
+ annotations: Optional[Dict] = None,
+ resources: Optional[Dict] = None,
+ affinity: Optional[Dict] = None,
+ config_file: Optional[str] = None,
+ node_selectors: Optional[Dict] = None,
+ image_pull_secrets: Optional[str] = None,
+ service_account_name: str = 'default',
+ is_delete_operator_pod: bool = False,
+ hostnetwork: bool = False,
+ tolerations: Optional[List] = None,
+ configmaps: Optional[List] = None,
+ security_context: Optional[Dict] = None,
+ pod_runtime_info_envs: Optional[List[PodRuntimeInfoEnv]] = 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,
**kwargs):
if kwargs.get('xcom_push') is not None:
raise AirflowException("'xcom_push' was deprecated, use 'do_xcom_push' instead")
- super(KubernetesPodOperator, self).__init__(*args, resources=None, **kwargs)
+ super().__init__(resources=None, **kwargs)
self.pod = None
self.do_xcom_push = do_xcom_push
@@ -211,7 +217,7 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
self.node_selectors = node_selectors or {}
self.annotations = annotations or {}
self.affinity = affinity or {}
- self.resources = self._set_resources(resources)
+ self.resources = self._set_resources(resources) # noqa
self.config_file = config_file
self.image_pull_secrets = image_pull_secrets
self.service_account_name = service_account_name
@@ -226,12 +232,13 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
self.full_pod_spec = full_pod_spec
self.init_containers = init_containers or []
self.log_events_on_failure = log_events_on_failure
- self.pod_template_file = pod_template_file
self.priority_class_name = priority_class_name
+ self.pod_template_file = pod_template_file
self.name = self._set_name(name)
+ self.client = None
@staticmethod
- def create_labels_for_pod(context):
+ def create_labels_for_pod(context) -> dict:
"""
Generate labels for the pod to track the pod in case of Operator crash
@@ -254,7 +261,7 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
labels[label_id] = safe_label
return labels
- def execute(self, context):
+ def execute(self, context) -> Optional[str]:
try:
if self.in_cluster is not None:
client = kube_client.get_kube_client(in_cluster=self.in_cluster,
@@ -264,12 +271,11 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
client = kube_client.get_kube_client(cluster_context=self.cluster_context,
config_file=self.config_file)
+ self.client = client
+
# Add combination of labels to uniquely identify a running pod
labels = self.create_labels_for_pod(context)
- self.pod = self.create_pod_request_obj()
- self.namespace = self.pod.metadata.namespace
-
label_selector = self._get_pod_identifying_label_string(labels)
pod_list = client.list_namespaced_pod(self.namespace, label_selector=label_selector)
@@ -285,6 +291,7 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
try_numbers_match = self._try_numbers_match(context, pod_list.items[0])
final_state, result = self.handle_pod_overlap(labels, try_numbers_match, launcher, pod_list)
else:
+ self.log.info("creating pod with labels %s and launcher %s", labels, launcher)
final_state, _, result = self.create_new_pod_for_operator(labels, launcher)
if final_state != State.SUCCESS:
raise AirflowException(
@@ -295,9 +302,11 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
def handle_pod_overlap(self, labels, try_numbers_match, launcher, pod_list):
"""
- In cases where the Scheduler restarts while a KubernetsPodOperator task is running,
+
+ In cases where the Scheduler restarts while a KubernetesPodOperator task is running,
this function will either continue to monitor the existing pod or launch a new pod
based on the `reattach_on_restart` parameter.
+
:param labels: labels used to determine if a pod is repeated
:type labels: dict
:param try_numbers_match: do the try numbers match? Only needed for logging purposes
@@ -311,11 +320,11 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
log_line = "found a running pod with labels {} but a different try_number.".format(labels)
if self.reattach_on_restart:
- log_line = log_line + " Will attach to this pod and monitor instead of starting new one"
+ log_line += " Will attach to this pod and monitor instead of starting new one"
self.log.info(log_line)
final_state, result = self.monitor_launched_pod(launcher, pod_list.items[0])
else:
- log_line = log_line + "creating pod with labels {} and launcher {}".format(labels, launcher)
+ log_line += "creating pod with labels {} and launcher {}".format(labels, launcher)
self.log.info(log_line)
final_state, _, result = self.create_new_pod_for_operator(labels, launcher)
return final_state, result
@@ -341,11 +350,24 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
validate_key(name, max_length=220)
return re.sub(r'[^a-z0-9.-]+', '-', name.lower())
- def create_pod_request_obj(self):
+ def create_new_pod_for_operator(self, labels, launcher) -> Tuple[State, k8s.V1Pod, Optional[str]]:
"""
- Creates a V1Pod based on user parameters. Note that a `pod` or `pod_template_file`
- will supersede all other values.
+ Creates a new pod and monitors for duration of task
+
+ :param labels: labels used to track pod
+ :param launcher: pod launcher that will manage launching and monitoring pods
+ :return:
"""
+ if not (self.full_pod_spec or self.pod_template_file):
+ # Add Airflow Version to the label
+ # And a label to identify that pod is launched by KubernetesPodOperator
+ self.labels.update(
+ {
+ 'airflow_version': airflow_version.replace('+', '-'),
+ 'kubernetes_pod_operator': 'True',
+ }
+ )
+ self.labels.update(labels)
pod = pod_generator.PodGenerator(
image=self.image,
namespace=self.namespace,
@@ -366,68 +388,46 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
configmaps=self.configmaps,
security_context=self.security_context,
dnspolicy=self.dnspolicy,
+ schedulername=self.schedulername,
init_containers=self.init_containers,
restart_policy='Never',
- schedulername=self.schedulername,
- pod_template_file=self.pod_template_file,
priority_class_name=self.priority_class_name,
+ pod_template_file=self.pod_template_file,
pod=self.full_pod_spec,
).gen_pod()
- # noinspection PyTypeChecker
pod = append_to_pod(
pod,
- self.pod_runtime_info_envs + # type: ignore
+ self.pod_runtime_info_envs +
self.ports + # type: ignore
- self.resources + # type: ignore
+ self.resources +
self.secrets + # type: ignore
self.volumes + # type: ignore
self.volume_mounts # type: ignore
)
- return pod
-
- def create_new_pod_for_operator(self, labels, launcher):
- """
- Creates a new pod and monitors for duration of task
-
- @param labels: labels used to track pod
- @param launcher: pod launcher that will manage launching and monitoring pods
- @return:
- """
- if not (self.full_pod_spec or self.pod_template_file):
- # Add Airflow Version to the label
- # And a label to identify that pod is launched by KubernetesPodOperator
- self.labels.update(
- {
- 'airflow_version': airflow_version.replace('+', '-'),
- 'kubernetes_pod_operator': 'True',
- }
- )
- self.labels.update(labels)
- self.pod.metadata.labels = self.labels
- self.log.debug("Starting pod:\n%s", yaml.safe_dump(self.pod.to_dict()))
-
+ self.pod = pod
+ self.log.debug("Starting pod:\n%s", yaml.safe_dump(pod.to_dict()))
try:
launcher.start_pod(
- self.pod,
+ pod,
startup_timeout=self.startup_timeout_seconds)
- final_state, result = launcher.monitor_pod(pod=self.pod, get_logs=self.get_logs)
- except AirflowException as ex:
+ final_state, result = launcher.monitor_pod(pod=pod, get_logs=self.get_logs)
+ except AirflowException:
if self.log_events_on_failure:
- for event in launcher.read_pod_events(self.pod).items:
+ for event in launcher.read_pod_events(pod).items:
self.log.error("Pod Event: %s - %s", event.reason, event.message)
- raise AirflowException('Pod Launching failed: {error}'.format(error=ex))
+ raise
finally:
if self.is_delete_operator_pod:
- launcher.delete_pod(self.pod)
- return final_state, self.pod, result
+ launcher.delete_pod(pod)
+ return final_state, pod, result
- def monitor_launched_pod(self, launcher, pod):
+ def monitor_launched_pod(self, launcher, pod) -> Tuple[State, Optional[str]]:
"""
- Montitors a pod to completion that was created by a previous KubernetesPodOperator
+ Monitors a pod to completion that was created by a previous KubernetesPodOperator
- @param launcher: pod launcher that will manage launching and monitoring pods
+ :param launcher: pod launcher that will manage launching and monitoring pods
:param pod: podspec used to find pod using k8s API
:return:
"""
@@ -444,3 +444,10 @@ class KubernetesPodOperator(BaseOperator): # pylint: disable=too-many-instance-
'Pod returned a failure: {state}'.format(state=final_state)
)
return final_state, result
+
+ def on_kill(self) -> None:
+ if self.pod:
+ pod: k8s.V1Pod = self.pod
+ namespace = pod.metadata.namespace
+ name = pod.metadata.name
+ self.client.delete_namespaced_pod(name=name, namespace=namespace, grace_period_seconds=0)
diff --git a/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes_tests/test_kubernetes_pod_operator.py
index eabcc02..efc48b8 100644
--- a/kubernetes_tests/test_kubernetes_pod_operator.py
+++ b/kubernetes_tests/test_kubernetes_pod_operator.py
@@ -1010,5 +1010,33 @@ class TestKubernetesPodOperatorSystem(unittest.TestCase):
do_xcom_push=False,
)
+ @mock.patch("airflow.kubernetes.pod_launcher.PodLauncher.monitor_pod")
+ def test_on_kill(self,
+ monitor_mock): # pylint: disable=unused-argument
+ from airflow.utils.state import State
+ client = kube_client.get_kube_client(in_cluster=False)
+ name = "test"
+ namespace = "default"
+ k = KubernetesPodOperator(
+ namespace='default',
+ image="ubuntu:16.04",
+ cmds=["bash", "-cx"],
+ arguments=["sleep 1000"],
+ labels={"foo": "bar"},
+ name="test",
+ task_id=name,
+ in_cluster=False,
+ do_xcom_push=False,
+ )
+ context = create_context(k)
+ monitor_mock.return_value = (State.SUCCESS, None)
+ k.execute(context)
+ name = k.pod.metadata.name
+ pod = client.read_namespaced_pod(name=name, namespace=namespace)
+ self.assertEqual(pod.status.phase, "Running")
+ k.on_kill()
+ with self.assertRaises(ApiException):
+ # pod should be deleted
+ client.read_namespaced_pod(name=name, namespace=namespace)
# pylint: enable=unused-argument