From cf3696fa02b381871479013b81076b7faceef5a3 Mon Sep 17 00:00:00 2001 From: Hossein Torabi Date: Mon, 10 Jul 2023 14:21:56 +0200 Subject: [PATCH] fix spark compability --- .../cncf/kubernetes/example_dags/__init__.py | 16 ++++++ .../spark-operator-with-sensor.py | 31 +++++++++++ .../example_dags/spark-operator-with-watch.py | 43 +++++++++++++++ .../kubernetes/example_dags/spark-pi.yaml | 55 +++++++++++++++++++ .../cncf/kubernetes/hooks/kubernetes.py | 23 +++++++- .../kubernetes/operators/spark_kubernetes.py | 7 +++ .../cncf/kubernetes/hooks/test_kubernetes.py | 31 +++++++++++ 7 files changed, 205 insertions(+), 1 deletion(-) create mode 100644 airflow/providers/cncf/kubernetes/example_dags/__init__.py create mode 100644 airflow/providers/cncf/kubernetes/example_dags/spark-operator-with-sensor.py create mode 100644 airflow/providers/cncf/kubernetes/example_dags/spark-operator-with-watch.py create mode 100644 airflow/providers/cncf/kubernetes/example_dags/spark-pi.yaml diff --git a/airflow/providers/cncf/kubernetes/example_dags/__init__.py b/airflow/providers/cncf/kubernetes/example_dags/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow/providers/cncf/kubernetes/example_dags/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/airflow/providers/cncf/kubernetes/example_dags/spark-operator-with-sensor.py b/airflow/providers/cncf/kubernetes/example_dags/spark-operator-with-sensor.py new file mode 100644 index 0000000000000..d7d47418c0646 --- /dev/null +++ b/airflow/providers/cncf/kubernetes/example_dags/spark-operator-with-sensor.py @@ -0,0 +1,31 @@ +# +# 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. +"""Example of a DAG that uses :py:class:`SparkKubernetesOperator` with `watch` that is responsible +for the spark job states and logs. +""" +from __future__ import annotations + +import datetime + +from airflow import DAG +from airflow.providers.cncf.kubernetes.operators.spark_kubernetes import SparkKubernetesOperator + +with DAG("example_spark_k8s_operator_with_watch", start_date=datetime.datetime(2023, 1, 1)) as dag: + SparkKubernetesOperator( + task_id="spark-operator", application_file="spark-pi.yaml", namespace="spark-apps", watch=True + ) diff --git a/airflow/providers/cncf/kubernetes/example_dags/spark-operator-with-watch.py b/airflow/providers/cncf/kubernetes/example_dags/spark-operator-with-watch.py new file mode 100644 index 0000000000000..9d0ed4f9d2888 --- /dev/null +++ b/airflow/providers/cncf/kubernetes/example_dags/spark-operator-with-watch.py @@ -0,0 +1,43 @@ +# +# 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. +"""Example of a DAG that uses :py:class:`SparkKubernetesOperator` with sensor. +The operator creates the spark job on the kubernetes and the sensor watches the state of the job. +""" +from __future__ import annotations + +import datetime + +from airflow import DAG +from airflow.providers.cncf.kubernetes.operators.spark_kubernetes import SparkKubernetesOperator +from airflow.providers.cncf.kubernetes.sensors.spark_kubernetes import SparkKubernetesSensor + +with DAG("example_spark_k8s_with_sensor", start_date=datetime.datetime(2023, 1, 1)) as dag: + t1 = SparkKubernetesOperator( + task_id="spark_pi_submit", + namespace="default", + application_file="spark-pi.yaml", + do_xcom_push=True, + ) + + t2 = SparkKubernetesSensor( + task_id="spark_pi_sensor", + namespace="default", + application_name="{{ task_instance.xcom_pull(task_ids='spark_pi_submit')['metadata']['name'] }}", + ) + + t1 >> t2 diff --git a/airflow/providers/cncf/kubernetes/example_dags/spark-pi.yaml b/airflow/providers/cncf/kubernetes/example_dags/spark-pi.yaml new file mode 100644 index 0000000000000..412a2c78be9d2 --- /dev/null +++ b/airflow/providers/cncf/kubernetes/example_dags/spark-pi.yaml @@ -0,0 +1,55 @@ +# 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. +--- +apiVersion: "sparkoperator.k8s.io/v1beta2" +kind: SparkApplication +metadata: + name: spark-pi +spec: + type: Scala + mode: cluster + image: "gcr.io/spark-operator/spark:v3.1.1" + imagePullPolicy: Always + mainClass: org.apache.spark.examples.SparkPi + mainApplicationFile: "local:///opt/spark/examples/jars/spark-examples_2.12-3.1.1.jar" + sparkVersion: "3.1.1" + restartPolicy: + type: Never + volumes: + - name: "test-volume" + hostPath: + path: "/tmp" + type: Directory + driver: + cores: 1 + coreLimit: "1200m" + memory: "512m" + labels: + version: 3.1.1 + serviceAccount: spark + volumeMounts: + - name: "test-volume" + mountPath: "/tmp" + executor: + cores: 1 + instances: 1 + memory: "512m" + labels: + version: 3.1.1 + volumeMounts: + - name: "test-volume" + mountPath: "/tmp" diff --git a/airflow/providers/cncf/kubernetes/hooks/kubernetes.py b/airflow/providers/cncf/kubernetes/hooks/kubernetes.py index f237461155dfb..ca91dc663699f 100644 --- a/airflow/providers/cncf/kubernetes/hooks/kubernetes.py +++ b/airflow/providers/cncf/kubernetes/hooks/kubernetes.py @@ -277,7 +277,13 @@ def custom_object_client(self) -> client.CustomObjectsApi: return client.CustomObjectsApi(api_client=self.api_client) def create_custom_object( - self, group: str, version: str, plural: str, body: str | dict, namespace: str | None = None + self, + group: str, + version: str, + plural: str, + body: str | dict, + namespace: str | None = None, + delete_existing_object: bool = True, ): """ Creates custom resource definition object in Kubernetes. @@ -287,6 +293,7 @@ def create_custom_object( :param plural: api plural :param body: crd object definition :param namespace: kubernetes namespace + :param delete_existing_object: whether to delete the existing object before creating or not """ api: client.CustomObjectsApi = self.custom_object_client @@ -295,6 +302,20 @@ def create_custom_object( else: body_dict = body + if delete_existing_object and "name" in body_dict["metadata"]: + try: + api.delete_namespaced_custom_object( + group=group, + version=version, + namespace=namespace, + plural=plural, + name=body_dict["metadata"]["name"], + ) + + self.log.warning("Deleted %s with the same name", plural) + except client.ApiException: + self.log.info("%s %s not found", plural, body_dict["metadata"]["name"]) + response = api.create_namespaced_custom_object( group=group, version=version, diff --git a/airflow/providers/cncf/kubernetes/operators/spark_kubernetes.py b/airflow/providers/cncf/kubernetes/operators/spark_kubernetes.py index ccc13469faaeb..b3e3e00b9e7b1 100644 --- a/airflow/providers/cncf/kubernetes/operators/spark_kubernetes.py +++ b/airflow/providers/cncf/kubernetes/operators/spark_kubernetes.py @@ -46,6 +46,7 @@ class SparkKubernetesOperator(BaseOperator): :param api_group: kubernetes api group of sparkApplication :param api_version: kubernetes api version of sparkApplication :param watch: whether to watch the job status and logs or not + :param delete_existing_object: whether to delete the existing object before creating or not """ template_fields: Sequence[str] = ("application_file", "namespace") @@ -64,6 +65,7 @@ def __init__( cluster_context: str | None = None, config_file: str | None = None, watch: bool = False, + delete_existing_object: bool = True, **kwargs, ) -> None: super().__init__(**kwargs) @@ -77,6 +79,7 @@ def __init__( self.cluster_context = cluster_context self.config_file = config_file self.watch = watch + self.delete_existing_object = delete_existing_object self.hook = KubernetesHook( conn_id=self.kubernetes_conn_id, @@ -109,6 +112,7 @@ def execute(self, context: Context): namespace=namespace, ) is_job_created = True + is_deleted = self.delete_existing_object for event in namespace_event_stream: obj = event["object"] if event["object"].last_timestamp >= datetime.datetime.strptime( @@ -129,6 +133,9 @@ def execute(self, context: Context): "SparkApplicationFailed", "SparkApplicationDeleted", ]: + if is_deleted and obj.reason == "SparkApplicationDeleted": + is_deleted = False + continue is_job_created = False raise AirflowException(obj.message) elif obj.reason == "SparkApplicationCompleted": diff --git a/tests/providers/cncf/kubernetes/hooks/test_kubernetes.py b/tests/providers/cncf/kubernetes/hooks/test_kubernetes.py index ba151efaf7fa6..c15b3fef09a98 100644 --- a/tests/providers/cncf/kubernetes/hooks/test_kubernetes.py +++ b/tests/providers/cncf/kubernetes/hooks/test_kubernetes.py @@ -373,6 +373,37 @@ def test_missing_default_connection_is_ok(self, remove_default_conn): with pytest.raises(AirflowNotFoundException, match="The conn_id `some_conn` isn't defined"): hook.conn_extras + @patch("kubernetes.config.kube_config.KubeConfigLoader") + @patch("kubernetes.config.kube_config.KubeConfigMerger") + @patch(f"{HOOK_MODULE}.client.CustomObjectsApi") + def test_create_custom_object( + self, mock_custom_object_api, mock_kube_config_merger, mock_kube_config_loader + ): + hook = KubernetesHook() + hook.create_custom_object( + group="group", + version="version", + plural="plural", + namespace="namespace", + body={"metadata": {"name": "name"}}, + ) + + mock_custom_object_api.return_value.delete_namespaced_custom_object.assert_called_once_with( + group="group", + version="version", + plural="plural", + name="name", + namespace="namespace", + ) + + mock_custom_object_api.return_value.create_namespaced_custom_object.assert_called_once_with( + group="group", + version="version", + plural="plural", + namespace="namespace", + body={"metadata": {"name": "name"}}, + ) + @patch("kubernetes.config.kube_config.KubeConfigLoader") @patch("kubernetes.config.kube_config.KubeConfigMerger") @patch(f"{HOOK_MODULE}.client.CustomObjectsApi")