apache-airflow-providers-cncf-kubernetes 10.1.0rc2__py3-none-any.whl → 10.3.0__py3-none-any.whl

airflow/providers/cncf/kubernetes/get_provider_info.py

@@ -15,8 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE
-# OVERWRITTEN WHEN PREPARING PACKAGES.
+# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN!
 #
 # IF YOU WANT TO MODIFY THIS FILE, YOU SHOULD MODIFY THE TEMPLATE
 # `get_provider_info_TEMPLATE.py.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY
@@ -28,8 +27,9 @@ def get_provider_info():
         "name": "Kubernetes",
         "description": "`Kubernetes <https://kubernetes.io/>`__\n",
         "state": "ready",
-        "source-date-epoch": 1734537609,
+        "source-date-epoch": 1739959070,
         "versions": [
+            "10.3.0",
             "10.1.0",
             "10.0.1",
             "10.0.0",
@@ -100,27 +100,18 @@ def get_provider_info():
             "1.0.1",
             "1.0.0",
         ],
-        "dependencies": [
-            "aiofiles>=23.2.0",
-            "apache-airflow>=2.9.0",
-            "asgiref>=3.5.2",
-            "cryptography>=41.0.0",
-            "kubernetes>=29.0.0,<=31.0.0",
-            "kubernetes_asyncio>=29.0.0,<=31.0.0",
-            "google-re2>=1.0",
-        ],
         "integrations": [
             {
                 "integration-name": "Kubernetes",
                 "external-doc-url": "https://kubernetes.io/",
                 "how-to-guide": ["/docs/apache-airflow-providers-cncf-kubernetes/operators.rst"],
-                "logo": "/integration-logos/kubernetes/Kubernetes.png",
+                "logo": "/docs/integration-logos/Kubernetes.png",
                 "tags": ["software"],
             },
             {
                 "integration-name": "Spark on Kubernetes",
                 "external-doc-url": "https://github.com/GoogleCloudPlatform/spark-on-k8s-operator",
-                "logo": "/integration-logos/kubernetes/Spark-On-Kubernetes.png",
+                "logo": "/docs/integration-logos/Spark-On-Kubernetes.png",
                 "tags": ["software"],
             },
         ],
@@ -341,13 +332,6 @@ def get_provider_info():
                         "example": None,
                         "default": "True",
                     },
-                    "worker_pods_queued_check_interval": {
-                        "description": 'How often in seconds to check for task instances stuck in "queued" status without a pod\n',
-                        "version_added": None,
-                        "type": "integer",
-                        "example": None,
-                        "default": "60",
-                    },
                     "ssl_ca_cert": {
                         "description": "Path to a CA certificate to be used by the Kubernetes client to verify the server's SSL certificate.\n",
                         "version_added": None,
@@ -366,4 +350,13 @@ def get_provider_info():
             },
         },
         "executors": ["airflow.providers.cncf.kubernetes.kubernetes_executor.KubernetesExecutor"],
+        "dependencies": [
+            "aiofiles>=23.2.0",
+            "apache-airflow>=2.9.0",
+            "asgiref>=3.5.2",
+            "cryptography>=41.0.0",
+            "kubernetes>=29.0.0,<=31.0.0",
+            "kubernetes_asyncio>=29.0.0,<=31.0.0",
+            "google-re2>=1.0",
+        ],
     }

airflow/providers/cncf/kubernetes/hooks/kubernetes.py

@@ -29,9 +29,6 @@ import aiofiles
 import requests
 import tenacity
 from asgiref.sync import sync_to_async
-from kubernetes import client, config, utils, watch
-from kubernetes.client.models import V1Deployment
-from kubernetes.config import ConfigException
 from kubernetes_asyncio import client as async_client, config as async_config
 from urllib3.exceptions import HTTPError
 
@@ -46,6 +43,9 @@ from airflow.providers.cncf.kubernetes.utils.pod_manager import (
     container_is_running,
 )
 from airflow.utils import yaml
+from kubernetes import client, config, utils, watch
+from kubernetes.client.models import V1Deployment
+from kubernetes.config import ConfigException
 
 if TYPE_CHECKING:
     from kubernetes.client import V1JobList
@@ -734,9 +734,19 @@ class AsyncKubernetesHook(KubernetesHook):
         """Return Kubernetes API session for use with requests."""
         in_cluster = self._coalesce_param(self.in_cluster, await self._get_field("in_cluster"))
         cluster_context = self._coalesce_param(self.cluster_context, await self._get_field("cluster_context"))
+        kubeconfig_path = await self._get_field("kube_config_path")
         kubeconfig = await self._get_field("kube_config")
+        num_selected_configuration = sum(
+            1 for o in [in_cluster, kubeconfig, kubeconfig_path, self.config_dict] if o
+        )
 
-        num_selected_configuration = sum(1 for o in [in_cluster, kubeconfig, self.config_dict] if o)
+        async def api_client_from_kubeconfig_file(_kubeconfig_path: str | None):
+            await async_config.load_kube_config(
+                config_file=_kubeconfig_path,
+                client_configuration=self.client_configuration,
+                context=cluster_context,
+            )
+            return async_client.ApiClient()
 
         if num_selected_configuration > 1:
             raise AirflowException(
@@ -757,6 +767,11 @@ class AsyncKubernetesHook(KubernetesHook):
             await async_config.load_kube_config_from_dict(self.config_dict)
             return async_client.ApiClient()
 
+        if kubeconfig_path is not None:
+            self.log.debug("loading kube_config from: %s", kubeconfig_path)
+            self._is_in_cluster = False
+            return await api_client_from_kubeconfig_file(kubeconfig_path)
+
         if kubeconfig is not None:
             async with aiofiles.tempfile.NamedTemporaryFile() as temp_config:
                 self.log.debug(
@@ -766,12 +781,7 @@ class AsyncKubernetesHook(KubernetesHook):
                 await temp_config.write(kubeconfig.encode())
                 await temp_config.flush()
                 self._is_in_cluster = False
-                await async_config.load_kube_config(
-                    config_file=temp_config.name,
-                    client_configuration=self.client_configuration,
-                    context=cluster_context,
-                )
-            return async_client.ApiClient()
+                return await api_client_from_kubeconfig_file(temp_config.name)
         self.log.debug(LOADING_KUBE_CONFIG_FILE_RESOURCE.format("default configuration file"))
         await async_config.load_kube_config(
             client_configuration=self.client_configuration,

airflow/providers/cncf/kubernetes/kube_config.py

@@ -76,10 +76,6 @@ class KubeConfig:
         # interact with cluster components.
         self.executor_namespace = conf.get(self.kubernetes_section, "namespace")
 
-        self.worker_pods_queued_check_interval = conf.getint(
-            self.kubernetes_section, "worker_pods_queued_check_interval"
-        )
-
         self.kube_client_request_args = conf.getjson(
             self.kubernetes_section, "kube_client_request_args", fallback={}
         )

airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py

@@ -23,11 +23,11 @@ from functools import cache
 from typing import TYPE_CHECKING
 
 import pendulum
-from kubernetes.client.rest import ApiException
 from slugify import slugify
 
 from airflow.configuration import conf
 from airflow.providers.cncf.kubernetes.backcompat import get_logical_date_key
+from kubernetes.client.rest import ApiException
 
 if TYPE_CHECKING:
     from airflow.models.taskinstancekey import TaskInstanceKey

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

@@ -24,8 +24,6 @@ from datetime import datetime as dt
 from functools import cached_property
 
 import tenacity
-from kubernetes.client import CoreV1Api, CustomObjectsApi, models as k8s
-from kubernetes.client.rest import ApiException
 
 from airflow.exceptions import AirflowException
 from airflow.providers.cncf.kubernetes.resource_convert.configmap import (
@@ -39,6 +37,8 @@ from airflow.providers.cncf.kubernetes.resource_convert.secret import (
 )
 from airflow.providers.cncf.kubernetes.utils.pod_manager import PodManager
 from airflow.utils.log.logging_mixin import LoggingMixin
+from kubernetes.client import CoreV1Api, CustomObjectsApi, models as k8s
+from kubernetes.client.rest import ApiException
 
 
 def should_retry_start_spark_job(exception: BaseException) -> bool:
@@ -291,7 +291,7 @@ class CustomObjectLauncher(LoggingMixin):
             # Wait for the driver pod to come alive
             self.pod_spec = k8s.V1Pod(
                 metadata=k8s.V1ObjectMeta(
-                    labels=self.spark_obj_spec["spec"]["driver"]["labels"],
+                    labels=self.spark_obj_spec["spec"]["driver"].get("labels"),
                     name=self.spark_obj_spec["metadata"]["name"] + "-driver",
                     namespace=self.namespace,
                 )

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

@@ -26,10 +26,6 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from kubernetes.client import BatchV1Api, models as k8s
-from kubernetes.client.api_client import ApiClient
-from kubernetes.client.rest import ApiException
-
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
@@ -44,6 +40,9 @@ from airflow.providers.cncf.kubernetes.triggers.job import KubernetesJobTrigger
 from airflow.providers.cncf.kubernetes.utils.pod_manager import EMPTY_XCOM_RESULT, PodNotFoundException
 from airflow.utils import yaml
 from airflow.utils.context import Context
+from kubernetes.client import BatchV1Api, models as k8s
+from kubernetes.client.api_client import ApiClient
+from kubernetes.client.rest import ApiException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -167,6 +166,7 @@ class KubernetesJobOperator(KubernetesPodOperator):
         ti.xcom_push(key="job_name", value=self.job.metadata.name)
         ti.xcom_push(key="job_namespace", value=self.job.metadata.namespace)
 
+        self.pod: k8s.V1Pod | None
         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,

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

@@ -22,12 +22,11 @@ import json
 from collections.abc import Sequence
 from functools import cached_property
 
-from kubernetes.utils import FailToCreateError
-
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
 from airflow.providers.cncf.kubernetes.operators.job import KubernetesJobOperator
+from kubernetes.utils import FailToCreateError
 
 
 class KubernetesInstallKueueOperator(BaseOperator):
@@ -95,6 +94,7 @@ class KubernetesStartKueueJobOperator(KubernetesJobOperator):
         super().__init__(*args, **kwargs)
         self.queue_name = queue_name
 
+        self.suspend: bool
         if self.suspend is False:
             raise AirflowException(
                 "The `suspend` parameter can't be False. If you want to use Kueue for running Job"

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

@@ -26,19 +26,16 @@ import os
 import re
 import shlex
 import string
-from collections.abc import Container, Iterable, Mapping, Sequence
+from collections.abc import Container, Iterable, Sequence
 from contextlib import AbstractContextManager
 from enum import Enum
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Callable, Literal
 
-import kubernetes
 import tenacity
-from kubernetes.client import CoreV1Api, V1Pod, models as k8s
-from kubernetes.client.exceptions import ApiException
-from kubernetes.stream import stream
 from urllib3.exceptions import HTTPError
 
+import kubernetes
 from airflow.configuration import conf
 from airflow.exceptions import (
     AirflowException,
@@ -84,13 +81,21 @@ from airflow.settings import pod_mutation_hook
 from airflow.utils import yaml
 from airflow.utils.helpers import prune_dict, validate_key
 from airflow.version import version as airflow_version
+from kubernetes.client import CoreV1Api, V1Pod, models as k8s
+from kubernetes.client.exceptions import ApiException
+from kubernetes.stream import stream
 
 if TYPE_CHECKING:
     import jinja2
     from pendulum import DateTime
 
     from airflow.providers.cncf.kubernetes.secret import Secret
-    from airflow.utils.context import Context
+
+    try:
+        from airflow.sdk.definitions.context import Context
+    except ImportError:
+        # TODO: Remove once provider drops support for Airflow 2
+        from airflow.utils.context import Context
 
 alphanum_lower = string.ascii_lowercase + string.digits
 
@@ -238,6 +243,8 @@ class KubernetesPodOperator(BaseOperator):
 
     template_fields: Sequence[str] = (
         "image",
+        "name",
+        "hostname",
         "cmds",
         "annotations",
         "arguments",
@@ -319,7 +326,9 @@ 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,
+        callbacks: (
+            list[type[KubernetesPodOperatorCallback]] | type[KubernetesPodOperatorCallback] | None
+        ) = None,
         progress_callback: Callable[[str], None] | None = None,
         logging_interval: int | None = None,
         **kwargs,
@@ -384,7 +393,7 @@ class KubernetesPodOperator(BaseOperator):
         self.priority_class_name = priority_class_name
         self.pod_template_file = pod_template_file
         self.pod_template_dict = pod_template_dict
-        self.name = self._set_name(name)
+        self.name = name
         self.random_name_suffix = random_name_suffix
         self.termination_grace_period = termination_grace_period
         self.pod_request_obj: k8s.V1Pod | None = None
@@ -410,7 +419,7 @@ 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.callbacks = [] if not callbacks else callbacks if isinstance(callbacks, list) else [callbacks]
         self._killed: bool = False
 
     @cached_property
@@ -423,7 +432,7 @@ class KubernetesPodOperator(BaseOperator):
     def _render_nested_template_fields(
         self,
         content: Any,
-        context: Mapping[str, Any],
+        context: Context,
         jinja_env: jinja2.Environment,
         seen_oids: set,
     ) -> None:
@@ -480,11 +489,11 @@ class KubernetesPodOperator(BaseOperator):
         }
 
         map_index = ti.map_index
-        if map_index >= 0:
-            labels["map_index"] = map_index
+        if map_index is not None and map_index >= 0:
+            labels["map_index"] = str(map_index)
 
         if include_try_number:
-            labels.update(try_number=ti.try_number)
+            labels.update(try_number=str(ti.try_number))
         # In the case of sub dags this is just useful
         # TODO: Remove this when the minimum version of Airflow is bumped to 3.0
         if getattr(context["dag"], "parent_dag", False):
@@ -514,8 +523,9 @@ class KubernetesPodOperator(BaseOperator):
     @cached_property
     def client(self) -> CoreV1Api:
         client = self.hook.core_v1_client
-        if self.callbacks:
-            self.callbacks.on_sync_client_creation(client=client)
+
+        for callback in self.callbacks:
+            callback.on_sync_client_creation(client=client, operator=self)
         return client
 
     def find_pod(self, namespace: str, context: Context, *, exclude_checked: bool = True) -> k8s.V1Pod | None:
@@ -579,6 +589,7 @@ class KubernetesPodOperator(BaseOperator):
 
     def execute(self, context: Context):
         """Based on the deferrable parameter runs the pod asynchronously or synchronously."""
+        self.name = self._set_name(self.name)
         if not self.deferrable:
             return self.execute_sync(context)
 
@@ -589,6 +600,14 @@ class KubernetesPodOperator(BaseOperator):
         try:
             if self.pod_request_obj is None:
                 self.pod_request_obj = self.build_pod_request_obj(context)
+            for callback in self.callbacks:
+                callback.on_pod_manifest_created(
+                    pod_request=self.pod_request_obj,
+                    client=self.client,
+                    mode=ExecutionMode.SYNC,
+                    context=context,
+                    operator=self,
+                )
             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,
@@ -601,28 +620,48 @@ 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
+            for callback in self.callbacks:
+                callback.on_pod_creation(
+                    pod=self.remote_pod,
+                    client=self.client,
+                    mode=ExecutionMode.SYNC,
+                    context=context,
+                    operator=self,
                 )
 
             self.await_init_containers_completion(pod=self.pod)
 
             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,
-                )
+                pod = self.find_pod(self.pod.metadata.namespace, context=context)
+                for callback in self.callbacks:
+                    callback.on_pod_starting(
+                        pod=pod,
+                        client=self.client,
+                        mode=ExecutionMode.SYNC,
+                        context=context,
+                        operator=self,
+                    )
 
             self.await_pod_completion(pod=self.pod)
             if self.callbacks:
-                self.callbacks.on_pod_completion(
-                    pod=self.find_pod(self.pod.metadata.namespace, context=context),
-                    client=self.client,
-                    mode=ExecutionMode.SYNC,
-                )
+                pod = self.find_pod(self.pod.metadata.namespace, context=context)
+                for callback in self.callbacks:
+                    callback.on_pod_completion(
+                        pod=pod,
+                        client=self.client,
+                        mode=ExecutionMode.SYNC,
+                        context=context,
+                        operator=self,
+                    )
+                for callback in self.callbacks:
+                    callback.on_pod_teardown(
+                        pod=pod,
+                        client=self.client,
+                        mode=ExecutionMode.SYNC,
+                        context=context,
+                        operator=self,
+                    )
 
             if self.do_xcom_push:
                 self.pod_manager.await_xcom_sidecar_container_start(pod=self.pod)
@@ -637,8 +676,14 @@ class KubernetesPodOperator(BaseOperator):
                 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)
+            for callback in self.callbacks:
+                callback.on_pod_cleanup(
+                    pod=pod_to_clean,
+                    client=self.client,
+                    mode=ExecutionMode.SYNC,
+                    context=context,
+                    operator=self,
+                )
 
         if self.do_xcom_push:
             return result
@@ -705,11 +750,15 @@ class KubernetesPodOperator(BaseOperator):
             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,
-            )
+            pod = self.find_pod(self.pod.metadata.namespace, context=context)
+            for callback in self.callbacks:
+                callback.on_pod_creation(
+                    pod=pod,
+                    client=self.client,
+                    mode=ExecutionMode.SYNC,
+                    context=context,
+                    operator=self,
+                )
         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)
@@ -770,10 +819,16 @@ class KubernetesPodOperator(BaseOperator):
             if not self.pod:
                 raise PodNotFoundException("Could not find pod after resuming from deferral")
 
-            if self.callbacks and event["status"] != "running":
-                self.callbacks.on_operator_resuming(
-                    pod=self.pod, event=event, client=self.client, mode=ExecutionMode.SYNC
-                )
+            if event["status"] != "running":
+                for callback in self.callbacks:
+                    callback.on_operator_resuming(
+                        pod=self.pod,
+                        event=event,
+                        client=self.client,
+                        mode=ExecutionMode.SYNC,
+                        context=context,
+                        operator=self,
+                    )
 
             follow = self.logging_interval is None
             last_log_time = event.get("last_log_time")
@@ -816,9 +871,9 @@ class KubernetesPodOperator(BaseOperator):
         except TaskDeferred:
             raise
         finally:
-            self._clean(event)
+            self._clean(event, context)
 
-    def _clean(self, event: dict[str, Any]) -> None:
+    def _clean(self, event: dict[str, Any], context: Context) -> None:
         if event["status"] == "running":
             return
         istio_enabled = self.is_istio_enabled(self.pod)
@@ -841,6 +896,7 @@ class KubernetesPodOperator(BaseOperator):
             self.post_complete_action(
                 pod=self.pod,
                 remote_pod=self.pod,
+                context=context,
             )
 
     def _write_logs(self, pod: k8s.V1Pod, follow: bool = False, since_time: DateTime | None = None) -> None:
@@ -870,14 +926,16 @@ class KubernetesPodOperator(BaseOperator):
                 e if not isinstance(e, ApiException) else e.reason,
             )
 
-    def post_complete_action(self, *, pod, remote_pod, **kwargs) -> None:
+    def post_complete_action(self, *, pod, remote_pod, context: Context, **kwargs) -> None:
         """Actions that must be done after operator finishes logic of the deferrable_execution."""
         self.cleanup(
             pod=pod,
             remote_pod=remote_pod,
         )
-        if self.callbacks:
-            self.callbacks.on_pod_cleanup(pod=pod, client=self.client, mode=ExecutionMode.SYNC)
+        for callback in self.callbacks:
+            callback.on_pod_cleanup(
+                pod=pod, client=self.client, mode=ExecutionMode.SYNC, operator=self, context=context
+            )
 
     def cleanup(self, pod: k8s.V1Pod, remote_pod: k8s.V1Pod):
         # Skip cleaning the pod in the following scenarios.

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

@@ -25,7 +25,6 @@ from typing import TYPE_CHECKING
 
 import tenacity
 import yaml
-from kubernetes.utils import create_from_yaml
 
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
@@ -33,6 +32,7 @@ from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
 from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import should_retry_creation
 from airflow.providers.cncf.kubernetes.utils.delete_from import delete_from_yaml
 from airflow.providers.cncf.kubernetes.utils.k8s_resource_iterator import k8s_resource_iterator
+from kubernetes.utils import create_from_yaml
 
 if TYPE_CHECKING:
     from kubernetes.client import ApiClient, CustomObjectsApi

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

@@ -17,12 +17,9 @@
 # under the License.
 from __future__ import annotations
 
-from collections.abc import Mapping
 from functools import cached_property
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
-
-from kubernetes.client import CoreV1Api, CustomObjectsApi, models as k8s
+from typing import TYPE_CHECKING, Any, cast
 
 from airflow.exceptions import AirflowException
 from airflow.providers.cncf.kubernetes import pod_generator
@@ -33,11 +30,16 @@ from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperato
 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
+from kubernetes.client import CoreV1Api, CustomObjectsApi, models as k8s
 
 if TYPE_CHECKING:
     import jinja2
 
-    from airflow.utils.context import Context
+    try:
+        from airflow.sdk.definitions.context import Context
+    except ImportError:
+        # TODO: Remove once provider drops support for Airflow 2
+        from airflow.utils.context import Context
 
 
 class SparkKubernetesOperator(KubernetesPodOperator):
@@ -114,6 +116,10 @@ class SparkKubernetesOperator(KubernetesPodOperator):
         self.success_run_history_limit = success_run_history_limit
         self.random_name_suffix = random_name_suffix
 
+        # fix mypy typing
+        self.base_container_name: str
+        self.container_logs: list[str]
+
         if self.base_container_name != self.BASE_CONTAINER_NAME:
             self.log.warning(
                 "base_container_name is not supported and will be overridden to %s", self.BASE_CONTAINER_NAME
@@ -129,7 +135,7 @@ class SparkKubernetesOperator(KubernetesPodOperator):
     def _render_nested_template_fields(
         self,
         content: Any,
-        context: Mapping[str, Any],
+        context: Context,
         jinja_env: jinja2.Environment,
         seen_oids: set,
     ) -> None:
@@ -174,12 +180,7 @@ class SparkKubernetesOperator(KubernetesPodOperator):
         return self._set_name(updated_name)
 
     @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:
+    def _get_ti_pod_labels(context: Context | None = None, include_try_number: bool = True) -> dict[str, str]:
         """
         Generate labels for the pod to track the pod in case of Operator crash.
 
@@ -190,8 +191,9 @@ class SparkKubernetesOperator(KubernetesPodOperator):
         if not context:
             return {}
 
-        ti = context["ti"]
-        run_id = context["run_id"]
+        context_dict = cast(dict, context)
+        ti = context_dict["ti"]
+        run_id = context_dict["run_id"]
 
         labels = {
             "dag_id": ti.dag_id,
@@ -210,8 +212,8 @@ class SparkKubernetesOperator(KubernetesPodOperator):
 
         # In the case of sub dags this is just useful
         # TODO: Remove this when the minimum version of Airflow is bumped to 3.0
-        if getattr(context["dag"], "is_subdag", False):
-            labels["parent_dag_id"] = context["dag"].parent_dag.dag_id
+        if getattr(context_dict["dag"], "is_subdag", False):
+            labels["parent_dag_id"] = context_dict["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():
@@ -232,9 +234,11 @@ class SparkKubernetesOperator(KubernetesPodOperator):
         """Templated body for CustomObjectLauncher."""
         return self.manage_template_specs()
 
-    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"
+    def find_spark_job(self, context, exclude_checked: bool = True):
+        label_selector = (
+            self._build_find_pod_label_selector(context, exclude_checked=exclude_checked)
+            + ",spark-role=driver"
+        )
         pod_list = self.client.list_namespaced_pod(self.namespace, label_selector=label_selector).items
 
         pod = None