apache-airflow-providers-cncf-kubernetes 7.13.0__py3-none-any.whl → 7.14.0rc1__py3-none-any.whl

airflow/providers/cncf/kubernetes/operators/pod.py

@@ -29,6 +29,7 @@ from contextlib import AbstractContextManager
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence
 
+import kubernetes
 from kubernetes.client import CoreV1Api, V1Pod, models as k8s
 from kubernetes.stream import stream
 from urllib3.exceptions import HTTPError
@@ -48,6 +49,7 @@ from airflow.providers.cncf.kubernetes.backcompat.backwards_compat_converters im
     convert_volume,
     convert_volume_mount,
 )
+from airflow.providers.cncf.kubernetes.callbacks import ExecutionMode, KubernetesPodOperatorCallback
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
 from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
     POD_NAME_MAX_LENGTH,
@@ -195,9 +197,12 @@ class KubernetesPodOperator(BaseOperator):
         Deprecated - use `on_finish_action` instead.
     :param termination_message_policy: The termination message policy of the base container.
         Default value is "File"
-    :param active_deadline_seconds: The active_deadline_seconds which matches to active_deadline_seconds
+    :param active_deadline_seconds: The active_deadline_seconds which translates to active_deadline_seconds
         in V1PodSpec.
+    :param callbacks: KubernetesPodOperatorCallback instance contains the callbacks methods on different step
+        of KubernetesPodOperator.
     :param progress_callback: Callback function for receiving k8s container logs.
+        `progress_callback` is deprecated, please use :param `callbacks` instead.
     """
 
     # !!! Changes in KubernetesPodOperator's arguments should be also reflected in !!!
@@ -225,6 +230,7 @@ class KubernetesPodOperator(BaseOperator):
         "volumes",
         "volume_mounts",
         "cluster_context",
+        "configmaps",
     )
     template_fields_renderers = {"env_vars": "py"}
 
@@ -289,6 +295,7 @@ class KubernetesPodOperator(BaseOperator):
         is_delete_operator_pod: None | bool = None,
         termination_message_policy: str = "File",
         active_deadline_seconds: int | None = None,
+        callbacks: type[KubernetesPodOperatorCallback] | None = None,
         progress_callback: Callable[[str], None] | None = None,
         **kwargs,
     ) -> None:
@@ -306,8 +313,9 @@ class KubernetesPodOperator(BaseOperator):
         if pod_runtime_info_envs:
             self.env_vars.extend([convert_pod_runtime_info_env(p) for p in pod_runtime_info_envs])
         self.env_from = env_from or []
-        if configmaps:
-            self.env_from.extend([convert_configmap(c) for c in configmaps])
+        self.configmaps = configmaps
+        if self.configmaps:
+            self.env_from.extend([convert_configmap(c) for c in self.configmaps])
         self.ports = [convert_port(p) for p in ports] if ports else []
         self.volume_mounts = [convert_volume_mount(v) for v in volume_mounts] if volume_mounts else []
         self.volumes = [convert_volume(volume) for volume in volumes] if volumes else []
@@ -380,6 +388,8 @@ class KubernetesPodOperator(BaseOperator):
 
         self._config_dict: dict | None = None  # TODO: remove it when removing convert_config_file_to_dict
         self._progress_callback = progress_callback
+        self.callbacks = callbacks
+        self._killed: bool = False
 
     @cached_property
     def _incluster_namespace(self):
@@ -403,11 +413,13 @@ class KubernetesPodOperator(BaseOperator):
             elif isinstance(content, k8s.V1ResourceRequirements):
                 template_fields = ("limits", "requests")
             elif isinstance(content, k8s.V1Volume):
-                template_fields = ("name", "persistent_volume_claim")
+                template_fields = ("name", "persistent_volume_claim", "config_map")
             elif isinstance(content, k8s.V1VolumeMount):
                 template_fields = ("name", "sub_path")
             elif isinstance(content, k8s.V1PersistentVolumeClaimVolumeSource):
                 template_fields = ("claim_name",)
+            elif isinstance(content, k8s.V1ConfigMapVolumeSource):
+                template_fields = ("name",)
             else:
                 template_fields = None
 
@@ -457,7 +469,9 @@ class KubernetesPodOperator(BaseOperator):
 
     @cached_property
     def pod_manager(self) -> PodManager:
-        return PodManager(kube_client=self.client, progress_callback=self._progress_callback)
+        return PodManager(
+            kube_client=self.client, callbacks=self.callbacks, progress_callback=self._progress_callback
+        )
 
     @cached_property
     def hook(self) -> PodOperatorHookProtocol:
@@ -471,7 +485,10 @@ class KubernetesPodOperator(BaseOperator):
 
     @cached_property
     def client(self) -> CoreV1Api:
-        return self.hook.core_v1_client
+        client = self.hook.core_v1_client
+        if self.callbacks:
+            self.callbacks.on_sync_client_creation(client=client)
+        return client
 
     def find_pod(self, namespace: str, context: Context, *, exclude_checked: bool = True) -> k8s.V1Pod | None:
         """Return an already-running pod for this task instance if one exists."""
@@ -536,11 +553,13 @@ class KubernetesPodOperator(BaseOperator):
     def execute_sync(self, context: Context):
         result = None
         try:
-            self.pod_request_obj = self.build_pod_request_obj(context)
-            self.pod = self.get_or_create_pod(  # must set `self.pod` for `on_kill`
-                pod_request_obj=self.pod_request_obj,
-                context=context,
-            )
+            if self.pod_request_obj is None:
+                self.pod_request_obj = self.build_pod_request_obj(context)
+            if self.pod is None:
+                self.pod = self.get_or_create_pod(  # must set `self.pod` for `on_kill`
+                    pod_request_obj=self.pod_request_obj,
+                    context=context,
+                )
             # push to xcom now so that if there is an error we still have the values
             ti = context["ti"]
             ti.xcom_push(key="pod_name", value=self.pod.metadata.name)
@@ -548,7 +567,17 @@ class KubernetesPodOperator(BaseOperator):
 
             # get remote pod for use in cleanup methods
             self.remote_pod = self.find_pod(self.pod.metadata.namespace, context=context)
+            if self.callbacks:
+                self.callbacks.on_pod_creation(
+                    pod=self.remote_pod, client=self.client, mode=ExecutionMode.SYNC
+                )
             self.await_pod_start(pod=self.pod)
+            if self.callbacks:
+                self.callbacks.on_pod_starting(
+                    pod=self.find_pod(self.pod.metadata.namespace, context=context),
+                    client=self.client,
+                    mode=ExecutionMode.SYNC,
+                )
 
             if self.get_logs:
                 self.pod_manager.fetch_requested_container_logs(
@@ -562,6 +591,12 @@ class KubernetesPodOperator(BaseOperator):
                 self.pod_manager.await_container_completion(
                     pod=self.pod, container_name=self.base_container_name
                 )
+            if self.callbacks:
+                self.callbacks.on_pod_completion(
+                    pod=self.find_pod(self.pod.metadata.namespace, context=context),
+                    client=self.client,
+                    mode=ExecutionMode.SYNC,
+                )
 
             if self.do_xcom_push:
                 self.pod_manager.await_xcom_sidecar_container_start(pod=self.pod)
@@ -571,10 +606,14 @@ class KubernetesPodOperator(BaseOperator):
                 self.pod, istio_enabled, self.base_container_name
             )
         finally:
+            pod_to_clean = self.pod or self.pod_request_obj
             self.cleanup(
-                pod=self.pod or self.pod_request_obj,
+                pod=pod_to_clean,
                 remote_pod=self.remote_pod,
             )
+            if self.callbacks:
+                self.callbacks.on_pod_cleanup(pod=pod_to_clean, client=self.client, mode=ExecutionMode.SYNC)
+
         if self.do_xcom_push:
             return result
 
@@ -584,6 +623,12 @@ class KubernetesPodOperator(BaseOperator):
             pod_request_obj=self.pod_request_obj,
             context=context,
         )
+        if self.callbacks:
+            self.callbacks.on_pod_creation(
+                pod=self.find_pod(self.pod.metadata.namespace, context=context),
+                client=self.client,
+                mode=ExecutionMode.SYNC,
+            )
         ti = context["ti"]
         ti.xcom_push(key="pod_name", value=self.pod.metadata.name)
         ti.xcom_push(key="pod_namespace", value=self.pod.metadata.namespace)
@@ -620,6 +665,10 @@ class KubernetesPodOperator(BaseOperator):
                 event["name"],
                 event["namespace"],
             )
+            if self.callbacks:
+                self.callbacks.on_operator_resuming(
+                    pod=pod, event=event, client=self.client, mode=ExecutionMode.SYNC
+                )
             if event["status"] in ("error", "failed", "timeout"):
                 # fetch some logs when pod is failed
                 if self.get_logs:
@@ -672,8 +721,15 @@ class KubernetesPodOperator(BaseOperator):
             pod=pod,
             remote_pod=remote_pod,
         )
+        if self.callbacks:
+            self.callbacks.on_pod_cleanup(pod=pod, client=self.client, mode=ExecutionMode.SYNC)
 
     def cleanup(self, pod: k8s.V1Pod, remote_pod: k8s.V1Pod):
+        # If a task got marked as failed, "on_kill" method would be called and the pod will be cleaned up
+        # there. Cleaning it up again will raise an exception (which might cause retry).
+        if self._killed:
+            return
+
         istio_enabled = self.is_istio_enabled(remote_pod)
         pod_phase = remote_pod.status.phase if hasattr(remote_pod, "status") else None
 
@@ -816,6 +872,7 @@ class KubernetesPodOperator(BaseOperator):
             )
 
     def on_kill(self) -> None:
+        self._killed = True
         if self.pod:
             pod = self.pod
             kwargs = {
@@ -824,7 +881,11 @@ class KubernetesPodOperator(BaseOperator):
             }
             if self.termination_grace_period is not None:
                 kwargs.update(grace_period_seconds=self.termination_grace_period)
-            self.client.delete_namespaced_pod(**kwargs)
+
+            try:
+                self.client.delete_namespaced_pod(**kwargs)
+            except kubernetes.client.exceptions.ApiException:
+                self.log.exception("Unable to delete pod %s", self.pod.metadata.name)
 
     def build_pod_request_obj(self, context: Context | None = None) -> k8s.V1Pod:
         """

airflow/providers/cncf/kubernetes/operators/spark_kubernetes.py

@@ -17,179 +17,264 @@
 # under the License.
 from __future__ import annotations
 
-import datetime
+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 python dictionary.
+    :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 | dict,
-        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 process_pod_deletion(self, pod, *, reraise=True):
+        if pod is not None:
+            if self.delete_on_termination:
+                self.log.info("Deleting spark job: %s", pod.metadata.name.replace("-driver", ""))
+                self.launcher.delete_spark_job(pod.metadata.name.replace("-driver", ""))
+            else:
+                self.log.info("skipping deleting spark job: %s", pod.metadata.name)
 
     @cached_property
     def hook(self) -> KubernetesHook:
-        return KubernetesHook(
+        hook = KubernetesHook(
             conn_id=self.kubernetes_conn_id,
-            in_cluster=self.in_cluster,
-            config_file=self.config_file,
-            cluster_context=self.cluster_context,
+            in_cluster=self.in_cluster or self.template_body.get("kubernetes", {}).get("in_cluster", False),
+            config_file=self.config_file
+            or self.template_body.get("kubernetes", {}).get("kube_config_file", None),
+            cluster_context=self.cluster_context
+            or self.template_body.get("kubernetes", {}).get("cluster_context", None),
         )
+        return hook
 
-    def _get_namespace_event_stream(self, namespace, query_kwargs=None):
-        try:
-            return Watch().stream(
-                self.hook.core_v1_client.list_namespaced_event,
-                namespace=namespace,
-                watch=True,
-                **(query_kwargs or {}),
-            )
-        except ApiException as e:
-            if e.status == 410:  # Resource version is too old
-                events: CoreV1EventList = self.hook.core_v1_client.list_namespaced_event(
-                    namespace=namespace, watch=False
-                )
-                resource_version = events.metadata.resource_version
-                query_kwargs["resource_version"] = resource_version
-                return self._get_namespace_event_stream(namespace, query_kwargs)
-            else:
-                raise
+    @cached_property
+    def client(self) -> CoreV1Api:
+        return self.hook.core_v1_client
 
-    def execute(self, context: Context):
-        if isinstance(self.application_file, str):
-            body = _load_body_to_dict(self.application_file)
-        else:
-            body = self.application_file
-        name = body["metadata"]["name"]
-        namespace = self.namespace or self.hook.get_namespace()
-
-        response = None
-        is_job_created = False
-        if self.watch:
-            try:
-                namespace_event_stream = self._get_namespace_event_stream(
-                    namespace=namespace,
-                    query_kwargs={
-                        "field_selector": f"involvedObject.kind=SparkApplication,involvedObject.name={name}"
-                    },
-                )
-
-                response = self.hook.create_custom_object(
-                    group=self.api_group,
-                    version=self.api_version,
-                    plural=self.plural,
-                    body=body,
-                    namespace=namespace,
-                )
-                is_job_created = True
-                for event in namespace_event_stream:
-                    obj = event["object"]
-                    if event["object"].last_timestamp >= datetime.datetime.strptime(
-                        response["metadata"]["creationTimestamp"], "%Y-%m-%dT%H:%M:%S%z"
-                    ):
-                        self.log.info(obj.message)
-                        if obj.reason == "SparkDriverRunning":
-                            pod_log_stream = Watch().stream(
-                                self.hook.core_v1_client.read_namespaced_pod_log,
-                                name=f"{name}-driver",
-                                namespace=namespace,
-                                timestamps=True,
-                            )
-                            for line in pod_log_stream:
-                                self.log.info(line)
-                        elif obj.reason in [
-                            "SparkApplicationSubmissionFailed",
-                            "SparkApplicationFailed",
-                            "SparkApplicationDeleted",
-                        ]:
-                            is_job_created = False
-                            raise AirflowException(obj.message)
-                        elif obj.reason == "SparkApplicationCompleted":
-                            break
-                        else:
-                            continue
-            except Exception:
-                if is_job_created:
-                    self.on_kill()
-                raise
+    @cached_property
+    def custom_obj_api(self) -> CustomObjectsApi:
+        return CustomObjectsApi()
 
-        else:
-            response = self.hook.create_custom_object(
-                group=self.api_group,
-                version=self.api_version,
-                plural=self.plural,
-                body=body,
-                namespace=namespace,
-            )
+    def execute(self, context: Context):
+        self.log.info("Creating sparkApplication.")
+        self.launcher = CustomObjectLauncher(
+            name=self.name,
+            namespace=self.namespace,
+            kube_client=self.client,
+            custom_obj_api=self.custom_obj_api,
+            template_body=self.template_body,
+        )
+        self.pod = self.get_or_create_spark_crd(self.launcher, context)
+        self.BASE_CONTAINER_NAME = "spark-kubernetes-driver"
+        self.pod_request_obj = self.launcher.pod_spec
 
-        return response
+        return super().execute(context=context)
 
     def on_kill(self) -> None:
-        if isinstance(self.application_file, str):
-            body = _load_body_to_dict(self.application_file)
-        else:
-            body = self.application_file
-        name = body["metadata"]["name"]
-        namespace = self.namespace or self.hook.get_namespace()
-        self.hook.delete_custom_object(
-            group=self.api_group,
-            version=self.api_version,
-            plural=self.plural,
-            namespace=namespace,
-            name=name,
-        )
+        if self.launcher:
+            self.log.debug("Deleting spark job for task %s", self.task_id)
+            self.launcher.delete_spark_job()
+
+    def patch_already_checked(self, pod: k8s.V1Pod, *, reraise=True):
+        """Add an "already checked" annotation to ensure we don't reattach on retries."""
+        pod.metadata.labels["already_checked"] = "True"
+        body = PodGenerator.serialize_pod(pod)
+        self.client.patch_namespaced_pod(pod.metadata.name, pod.metadata.namespace, body)
+
+    def dry_run(self) -> None:
+        """Prints out the spark job that would be created by this operator."""
+        print(prune_dict(self.launcher.body, mode="strict"))