hamedhsn commented on code in PR #22253:
URL: https://github.com/apache/airflow/pull/22253#discussion_r1403317000
##########
airflow/providers/cncf/kubernetes/operators/spark_kubernetes.py:
##########
@@ -17,173 +17,271 @@
# under the License.
from __future__ import annotations
-import datetime
+import json
+import re
from functools import cached_property
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, Any
-from kubernetes.client import ApiException
-from kubernetes.watch import Watch
+from kubernetes.client import CoreV1Api, CustomObjectsApi, models as k8s
from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
+from airflow.providers.cncf.kubernetes import pod_generator
from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook,
_load_body_to_dict
+from airflow.providers.cncf.kubernetes.operators.custom_object_launcher import
CustomObjectLauncher
+from airflow.providers.cncf.kubernetes.operators.pod import
KubernetesPodOperator
+from airflow.providers.cncf.kubernetes.pod_generator import MAX_LABEL_LEN,
PodGenerator
+from airflow.providers.cncf.kubernetes.utils.pod_manager import PodManager
+from airflow.utils.helpers import prune_dict
if TYPE_CHECKING:
- from kubernetes.client.models import CoreV1EventList
+ import jinja2
from airflow.utils.context import Context
-class SparkKubernetesOperator(BaseOperator):
+class SparkKubernetesOperator(KubernetesPodOperator):
"""
Creates sparkApplication object in kubernetes cluster.
.. seealso::
For more detail about Spark Application Object have a look at the
reference:
-
https://github.com/GoogleCloudPlatform/spark-on-k8s-operator/blob/v1beta2-1.1.0-2.4.5/docs/api-docs.md#sparkapplication
+
https://github.com/GoogleCloudPlatform/spark-on-k8s-operator/blob/v1beta2-1.3.3-3.1.1/docs/api-docs.md#sparkapplication
- :param application_file: Defines Kubernetes 'custom_resource_definition'
of 'sparkApplication' as either a
- path to a '.yaml' file, '.json' file, YAML string or JSON string.
+ :param application_file: filepath to kubernetes custom_resource_definition
of sparkApplication
+ :param kubernetes_conn_id: the connection to Kubernetes cluster
+ :param image: Docker image you wish to launch. Defaults to hub.docker.com,
+ :param code_path: path to the spark code in image,
:param namespace: kubernetes namespace to put sparkApplication
- :param kubernetes_conn_id: The :ref:`kubernetes connection id
<howto/connection:kubernetes>`
- for the to Kubernetes cluster.
- :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 cluster_context: context of the cluster
+ :param application_file: yaml file if passed
+ :param get_logs: get the stdout of the container as logs of the tasks.
+ :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 success_run_history_limit: Number of past successful runs of the
application to keep.
+ :param delete_on_termination: 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 startup_timeout_seconds: timeout in seconds to startup the pod.
+ :param log_events_on_failure: Log the pod's events if a failure occurs
+ :param reattach_on_restart: if the scheduler dies while the pod is
running, reattach and monitor
"""
- template_fields: Sequence[str] = ("application_file", "namespace")
- template_ext: Sequence[str] = (".yaml", ".yml", ".json")
+ template_fields = ["application_file", "namespace", "template_spec"]
+ template_fields_renderers = {"template_spec": "py"}
+ template_ext = ("yaml", "yml", "json")
ui_color = "#f4a460"
def __init__(
self,
*,
- application_file: str,
- namespace: str | None = None,
+ image: str | None = None,
+ code_path: str | None = None,
+ namespace: str = "default",
+ name: str = "default",
+ application_file: str | None = None,
+ template_spec=None,
+ get_logs: bool = True,
+ do_xcom_push: bool = False,
+ success_run_history_limit: int = 1,
+ startup_timeout_seconds=600,
+ log_events_on_failure: bool = False,
+ reattach_on_restart: bool = True,
+ delete_on_termination: bool = True,
kubernetes_conn_id: str = "kubernetes_default",
- api_group: str = "sparkoperator.k8s.io",
- api_version: str = "v1beta2",
- in_cluster: bool | None = None,
- cluster_context: str | None = None,
- config_file: str | None = None,
- watch: bool = False,
**kwargs,
) -> None:
- super().__init__(**kwargs)
- self.namespace = namespace
- self.kubernetes_conn_id = kubernetes_conn_id
- self.api_group = api_group
- self.api_version = api_version
- self.plural = "sparkapplications"
+ if kwargs.get("xcom_push") is not None:
+ raise AirflowException("'xcom_push' was deprecated, use
'do_xcom_push' instead")
+ super().__init__(name=name, **kwargs)
+ self.image = image
+ self.code_path = code_path
self.application_file = application_file
- self.in_cluster = in_cluster
- self.cluster_context = cluster_context
- self.config_file = config_file
- self.watch = watch
+ self.template_spec = template_spec
+ self.name = self.create_job_name()
+ self.kubernetes_conn_id = kubernetes_conn_id
+ self.startup_timeout_seconds = startup_timeout_seconds
+ self.reattach_on_restart = reattach_on_restart
+ self.delete_on_termination = delete_on_termination
+ self.do_xcom_push = do_xcom_push
+ self.namespace = namespace
+ self.get_logs = get_logs
+ self.log_events_on_failure = log_events_on_failure
+ self.success_run_history_limit = success_run_history_limit
+ self.template_body = self.manage_template_specs()
+
+ def _render_nested_template_fields(
+ self,
+ content: Any,
+ context: Context,
+ jinja_env: jinja2.Environment,
+ seen_oids: set,
+ ) -> None:
+ if id(content) not in seen_oids and isinstance(content, k8s.V1EnvVar):
+ seen_oids.add(id(content))
+ self._do_render_template_fields(content, ("value", "name"),
context, jinja_env, seen_oids)
+ return
+
+ super()._render_nested_template_fields(content, context, jinja_env,
seen_oids)
+
+ def manage_template_specs(self):
+ if self.application_file:
+ template_body = _load_body_to_dict(open(self.application_file))
+ elif self.template_spec:
+ template_body = self.template_spec
+ else:
+ raise AirflowException("either application_file or template_spec
should be passed")
+ if "spark" not in template_body:
+ template_body = {"spark": template_body}
+ return template_body
+
+ def create_job_name(self):
+ initial_name =
PodGenerator.make_unique_pod_id(self.task_id)[:MAX_LABEL_LEN]
+ return re.sub(r"[^a-z0-9-]+", "-", initial_name.lower())
+
+ @staticmethod
+ def _get_pod_identifying_label_string(labels) -> str:
+ filtered_labels = {label_id: label for label_id, label in
labels.items() if label_id != "try_number"}
+ return ",".join([label_id + "=" + label for label_id, label in
sorted(filtered_labels.items())])
+
+ @staticmethod
+ def create_labels_for_pod(context: dict | None = None, include_try_number:
bool = True) -> dict:
+ """
+ Generate labels for the pod to track the pod in case of Operator crash.
+
+ :param include_try_number: add try number to labels
+ :param context: task context provided by airflow DAG
+ :return: dict.
+ """
+ if not context:
+ return {}
+
+ ti = context["ti"]
+ run_id = context["run_id"]
+
+ labels = {
+ "dag_id": ti.dag_id,
+ "task_id": ti.task_id,
+ "run_id": run_id,
+ "spark_kubernetes_operator": "True",
+ # 'execution_date': context['ts'],
+ # 'try_number': context['ti'].try_number,
+ }
+
+ # If running on Airflow 2.3+:
+ map_index = getattr(ti, "map_index", -1)
+ if map_index >= 0:
+ labels["map_index"] = map_index
+
+ if include_try_number:
+ labels.update(try_number=ti.try_number)
+
+ # In the case of sub dags this is just useful
+ if context["dag"].is_subdag:
+ labels["parent_dag_id"] = context["dag"].parent_dag.dag_id
+ # Ensure that label is valid for Kube,
+ # and if not truncate/remove invalid chars and replace with short hash.
+ for label_id, label in labels.items():
+ safe_label = pod_generator.make_safe_label_value(str(label))
+ labels[label_id] = safe_label
+ return labels
+
+ @cached_property
+ def pod_manager(self) -> PodManager:
+ return PodManager(kube_client=self.client)
+
+ @staticmethod
+ def _try_numbers_match(context, pod) -> bool:
+ return pod.metadata.labels["try_number"] == context["ti"].try_number
+
+ def find_spark_job(self, context):
+ labels = self.create_labels_for_pod(context, include_try_number=False)
+ label_selector = self._get_pod_identifying_label_string(labels) +
",spark-role=driver"
+ pod_list = self.client.list_namespaced_pod(self.namespace,
label_selector=label_selector).items
+
+ pod = None
+ if len(pod_list) > 1: # and self.reattach_on_restart:
+ raise AirflowException(f"More than one pod running with labels:
{label_selector}")
+ elif len(pod_list) == 1:
+ pod = pod_list[0]
+ self.log.info(
+ "Found matching driver pod %s with labels %s",
pod.metadata.name, pod.metadata.labels
+ )
+ self.log.info("`try_number` of task_instance: %s",
context["ti"].try_number)
+ self.log.info("`try_number` of pod: %s",
pod.metadata.labels["try_number"])
+ return pod
+
+ def get_or_create_spark_crd(self, launcher: CustomObjectLauncher, context)
-> k8s.V1Pod:
+ if self.reattach_on_restart:
+ driver_pod = self.find_spark_job(context)
+ if driver_pod:
+ return driver_pod
+
+ driver_pod, spark_obj_spec = launcher.start_spark_job(
+ image=self.image, code_path=self.code_path,
startup_timeout=self.startup_timeout_seconds
+ )
+ return driver_pod
+
+ def extract_xcom(self, pod):
+ """Retrieves xcom value and kills xcom sidecar container."""
+ result = self.pod_manager.extract_xcom(pod)
+ self.log.info("xcom result: \n%s", result)
Review Comment:
we actually don't need this function, inherited from k8s operator so that
could be used.
--
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]
