diff --git a/providers/apache/spark/docs/operators.rst b/providers/apache/spark/docs/operators.rst index 92792a2672526..5003ea2b8fba5 100644 --- a/providers/apache/spark/docs/operators.rst +++ b/providers/apache/spark/docs/operators.rst @@ -284,10 +284,42 @@ before polling begins, and a retry reconnects to that pod instead of submitting conflicts with the flag and a ``ValueError`` will be raised at task start. * The Airflow worker must be able to reach the Kubernetes API server and have permission to read and delete pods in the driver's namespace; otherwise pod tracking and cleanup will fail. -* Pod completion is detected from ``pod.status.phase``. If your driver pods have sidecar - containers (e.g. Istio injection enabled for the driver namespace), the pod phase may not - advance to ``Succeeded`` until all sidecars exit. In that case the poll loop will wait - indefinitely — set ``execution_timeout`` as a hard bound. +* Set ``durable=True`` (the default) to enable crash recovery: the driver pod name is + persisted to task state before polling begins, so a worker crash and retry reconnects to the + existing pod instead of submitting a fresh one. Set ``durable=False`` to always + submit a fresh driver on retry. + +**Sidecar containers and driver container identification** + +Completion is detected from the driver container's own exit code rather than from +``pod.status.phase`` alone as long as the driver container can be identified. By default +the driver container is identified by name, preferring a container with ``driver`` in +its name, then one with ``spark`` in its name, falling back to the pod's only container if there +is just one. If this heuristic doesn't match your setup, set ``kubernetes_driver_container_name`` to +the exact container name: + +.. code-block:: python + + run_spark = SparkSubmitOperator( + task_id="run_spark", + application="local:///opt/spark/examples/jars/spark-examples.jar", + conn_id="spark_k8s", + deploy_mode="cluster", + track_driver_via_k8s_api=True, + kubernetes_driver_container_name="spark-kubernetes-driver", + ) + +If ``kubernetes_driver_container_name`` doesn't match any container on the pod, the task fails +immediately with a ``ValueError`` rather than silently falling back to the heuristic. + +If the pod phase reports ``Failed`` but the driver container itself exited 0 (for example, a +sidecar crashed after the driver finished), the operator logs a warning and still treats the task +as succeeded. + +If the driver container could not be identified, ``pod.status.phase`` will be used to track +completion. This matters if your driver pods have sidecar containers: the pod +phase may not advance to ``Succeeded`` until every container exits. To avoid indefinite +waits, set ``execution_timeout`` as a hard bound. Set ``durable=False`` to always submit a fresh driver on retry. diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py index 87d838333baea..38e8c21628d86 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py @@ -143,6 +143,7 @@ class SparkSubmitHook(BaseHook, LoggingMixin): ``keytab`` and ``principal`` configured use ``requests-kerberos`` automatically. Defaults to ``None`` (no auth for non-Kerberos connections). + :param kubernetes_driver_container_name: Name of the Spark driver container used for identification to track status """ conn_name_attr = "conn_id" @@ -273,6 +274,7 @@ def __init__( track_driver_via_k8s_api: bool = False, yarn_track_via_rm_api: bool = False, yarn_rm_auth: AuthBase | None = None, + kubernetes_driver_container_name: str | None = None, ) -> None: super().__init__() self._conf = conf or {} @@ -302,7 +304,7 @@ def __init__( self._verbose = verbose self._submit_sp: Any | None = None self._yarn_application_id: str | None = None - self._kubernetes_driver_pod: str | None = None + self._kubernetes_driver_pod_name: str | None = None self._kubernetes_application_id: str | None = None self.spark_binary = spark_binary self._properties_file = properties_file @@ -334,6 +336,7 @@ def __init__( # `_track_yarn_application` does not re-fetch the Spark connection # (and re-hit any configured Secrets Backend) on every iteration. self._yarn_rm_base_url: str | None = None + self.kubernetes_driver_container_name = kubernetes_driver_container_name def _resolve_should_track_driver_status(self) -> bool: """ @@ -857,13 +860,13 @@ def _process_spark_submit_log(self, itr: Iterator[Any]) -> None: # "pod name: -driver" and "submission ID spark:-driver" match_driver_pod = re.search(r"\s*pod name: ((.+?)-([a-z0-9]+)-driver$)", line) if match_driver_pod: - self._kubernetes_driver_pod = match_driver_pod.group(1) - self.log.info("Identified spark driver pod: %s", self._kubernetes_driver_pod) - if not self._kubernetes_driver_pod: + self._kubernetes_driver_pod_name = match_driver_pod.group(1) + self.log.info("Identified spark driver pod: %s", self._kubernetes_driver_pod_name) + if not self._kubernetes_driver_pod_name: match_submission_id = re.search(r"submission ID spark:(.+?-driver)", line) if match_submission_id: - self._kubernetes_driver_pod = match_submission_id.group(1) - self.log.info("Identified spark driver pod: %s", self._kubernetes_driver_pod) + self._kubernetes_driver_pod_name = match_submission_id.group(1) + self.log.info("Identified spark driver pod: %s", self._kubernetes_driver_pod_name) match_application_id = re.search(r"\s*spark-app-selector -> (spark-([a-z0-9]+)), ", line) if match_application_id: @@ -1177,15 +1180,16 @@ def _start_driver_status_tracking(self) -> None: def _poll_k8s_driver_via_api(self) -> str | None: """ - Poll the K8s driver pod phase until it reaches a terminal state. + Poll the K8s driver container status or pod phase until it reaches a terminal state. Returns the terminal phase string (e.g. ``"Succeeded"``) on normal completion, or ``None`` if the pod vanished mid-poll (404 — likely deleted by ``on_kill``). Raises ``RuntimeError`` on failure phases or unrecoverable API errors. """ - pod_name = self._kubernetes_driver_pod + kubernetes_driver_pod_name = self._kubernetes_driver_pod_name + kubernetes_driver_container_name = self.kubernetes_driver_container_name namespace = self._connection["namespace"] - app_id = self._kubernetes_application_id or pod_name + app_id = self._kubernetes_application_id or kubernetes_driver_pod_name client = kube_client.get_kube_client() poll_interval = max(self._status_poll_interval, 20) @@ -1201,26 +1205,29 @@ def _poll_k8s_driver_via_api(self) -> str | None: consecutive_api_errors = 0 max_consecutive_api_errors = 3 consecutive_pending = 0 - pending_warn_threshold = 10 + consecutive_waiting = 0 + waiting_or_pending_warn_threshold = 10 + terminal_phase: str | None = None + driver_container_logged = False try: - if not pod_name: + if not kubernetes_driver_pod_name: raise ValueError("K8s driver pod name not set; cannot poll status.") while True: try: - pod = client.read_namespaced_pod(pod_name, namespace) + pod = client.read_namespaced_pod(kubernetes_driver_pod_name, namespace) consecutive_api_errors = 0 except kube_client.ApiException as e: if e.status == 404: self.log.info( "Driver pod %s not found (404); pod was likely deleted by on_kill. Exiting poll loop.", - pod_name, + kubernetes_driver_pod_name, ) return None consecutive_api_errors += 1 self.log.warning( "ApiException polling pod %s (%d/%d): %s", - pod_name, + kubernetes_driver_pod_name, consecutive_api_errors, max_consecutive_api_errors, e, @@ -1232,7 +1239,56 @@ def _poll_k8s_driver_via_api(self) -> str | None: ) from e time.sleep(poll_interval) continue - + driver_container = None + + for container in pod.spec.containers: + if kubernetes_driver_container_name: + if kubernetes_driver_container_name.lower() == container.name.lower(): + driver_container = container + break + continue + if "driver" in container.name.lower(): + driver_container = container + break + if "spark" in container.name.lower(): + driver_container = container + if len(pod.spec.containers) == 1: + driver_container = container + if kubernetes_driver_container_name and not driver_container: + with contextlib.suppress(Exception): + self._delete_driver_pod() + raise ValueError( + f"The driver container name provided does not match any of the containers in pod " + f"{kubernetes_driver_pod_name}. Deleted the driver pod; check " + f"kubernetes_driver_container_name for a typo." + ) + container_completed = False + if driver_container: + if not driver_container_logged: + self.log.info("%s has been identified as the driver container", driver_container.name) + driver_container_logged = True + for status in pod.status.container_statuses or []: + if status.name == driver_container.name: + if status.state and status.state.terminated: + driver_exit_code = status.state.terminated.exit_code + if driver_exit_code == 0: + container_completed = True + break + raise RuntimeError( + f"Spark application {app_id} failed.\nThe driver container exited with a non-zero status code.\nExit code: {driver_exit_code}\nReason: {status.state.terminated.reason}" + ) + if status.state and status.state.waiting: + consecutive_waiting += 1 + if consecutive_waiting == waiting_or_pending_warn_threshold: + self.log.warning( + "Driver container %s has been waiting for %d polls (~%ds); " + "it may be unschedulable. Continuing to wait — set execution_timeout to bound wait time.", + driver_container.name, + consecutive_waiting, + consecutive_waiting * poll_interval, + ) + else: + consecutive_waiting = 0 phase = pod.status.phase or "Initializing" self.log.info("Application status for %s (phase: %s)", app_id, phase) if phase == "Succeeded": @@ -1249,20 +1305,27 @@ def _poll_k8s_driver_via_api(self) -> str | None: ) terminal_phase = phase break - if phase == "Failed": + if phase == "Failed" and not container_completed: container_state = "" if pod.status.container_statuses: cs = pod.status.container_statuses[0] if cs.state and cs.state.terminated: container_state = f" exit_code={cs.state.terminated.exit_code} reason={cs.state.terminated.reason}" raise RuntimeError(f"Spark application {app_id} failed (phase=Failed{container_state})") + if phase == "Failed" and container_completed and driver_container: + self.log.warning( + "Driver pod %s reported phase=Failed, but driver container %s exited 0; " + "treating the Spark application as succeeded.", + kubernetes_driver_pod_name, + driver_container.name, + ) if phase == "Pending": consecutive_pending += 1 - if consecutive_pending == pending_warn_threshold: + if consecutive_pending == waiting_or_pending_warn_threshold: self.log.warning( "Driver pod %s has been Pending for %d polls (~%ds); " "it may be unschedulable. Continuing to wait — set execution_timeout to bound wait time.", - pod_name, + kubernetes_driver_pod_name, consecutive_pending, consecutive_pending * poll_interval, ) @@ -1280,6 +1343,11 @@ def _poll_k8s_driver_via_api(self) -> str | None: ) else: consecutive_unknown = 0 + if container_completed: + # Driver container exited 0 — the application succeeded even if the + # pod phase still reads "Running" at this poll. + terminal_phase = "Succeeded" + break time.sleep(poll_interval) # Pod deletion is best-effort cleanup. If it fails (e.g. already garbage collected or RBAC # denied), suppress the error so terminal_phase is still returned and the task @@ -1314,19 +1382,19 @@ def _delete_driver_pod(self) -> None: """Delete the Kubernetes driver pod, logging a warning on failure.""" import kubernetes - self.log.info("Deleting driver pod %s on Kubernetes", self._kubernetes_driver_pod) + self.log.info("Deleting driver pod %s on Kubernetes", self._kubernetes_driver_pod_name) try: client = kube_client.get_kube_client() client.delete_namespaced_pod( - self._kubernetes_driver_pod, + self._kubernetes_driver_pod_name, self._connection["namespace"], body=kubernetes.client.V1DeleteOptions(), pretty=True, ) - self.log.info("Deleted driver pod %s", self._kubernetes_driver_pod) + self.log.info("Deleted driver pod %s", self._kubernetes_driver_pod_name) except kube_client.ApiException: self.log.exception( - "Exception when attempting to delete driver pod %s", self._kubernetes_driver_pod + "Exception when attempting to delete driver pod %s", self._kubernetes_driver_pod_name ) def on_kill(self) -> None: @@ -1342,7 +1410,7 @@ def on_kill(self) -> None: "Spark driver %s killed with return code: %s", self._driver_id, driver_kill.wait() ) - if self._should_track_driver_via_k8s_api() and self._kubernetes_driver_pod: + if self._should_track_driver_via_k8s_api() and self._kubernetes_driver_pod_name: # spark-submit exits early under waitAppCompletion=false, so _submit_sp.poll() is # not None during the poll loop — the deletion block below is skipped on kill. self._delete_driver_pod() @@ -1377,7 +1445,7 @@ def on_kill(self) -> None: ) as yarn_kill: self.log.info("YARN app killed with return code: %s", yarn_kill.wait()) - if self._kubernetes_driver_pod: + if self._kubernetes_driver_pod_name: self._delete_driver_pod() # Opt-in REST kill path — uses the same RM endpoint as polling, no diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py index 1a9d2748a651d..d1c9002ccb227 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py @@ -110,7 +110,7 @@ class _KubernetesSparkSubmitBackend(_SparkSubmitDeploymentBackend): def submit_job(self, context: Context) -> str | None: self.hook._conf[_K8S_WAIT_APP_COMPLETION_CONF] = "false" self.hook.submit(self.operator.application) - pod_name = self.hook._kubernetes_driver_pod + pod_name = self.hook._kubernetes_driver_pod_name namespace = self.hook._connection["namespace"] if not pod_name: raise RuntimeError("spark-submit did not capture a K8s driver pod name") @@ -145,7 +145,7 @@ def is_job_succeeded(self, status: str) -> bool: def poll_until_complete(self, external_id: str, context: Context) -> None: if external_id is not None: _, pod_name = self.operator._parse_k8s_external_id(external_id) - self.hook._kubernetes_driver_pod = pod_name + self.hook._kubernetes_driver_pod_name = pod_name terminal_phase = self.hook._poll_k8s_driver_via_api() # Cache only when the pod actually reached Succeeded, the 404/vanished path # returns None for cases like: pod deleted by on_kill or garbage collected after failure) @@ -341,6 +341,7 @@ class SparkSubmitOperator(ResumableJobMixin, BaseOperator): Requires Airflow 3.3 or newer; below that, ``durable`` has no effect -- setting it explicitly only emits a warning. :param reconnect_on_retry: deprecated, use ``durable`` instead. + :param kubernetes_driver_container_name: Name of the Spark driver container used for identification to track status """ # Generic key used across all Spark deployment modes (standalone driver ID, @@ -369,6 +370,7 @@ class SparkSubmitOperator(ResumableJobMixin, BaseOperator): "env_vars", "post_submit_commands", "properties_file", + "kubernetes_driver_container_name", ) def __init__( @@ -416,6 +418,7 @@ def __init__( ), reconnect_on_retry: bool | None = None, durable: bool | None = None, + kubernetes_driver_container_name: str | None = None, **kwargs: Any, ) -> None: if reconnect_on_retry is not None: @@ -471,6 +474,7 @@ def __init__( self._track_driver_via_k8s_api = track_driver_via_k8s_api self._openlineage_inject_parent_job_info = openlineage_inject_parent_job_info self._openlineage_inject_transport_info = openlineage_inject_transport_info + self.kubernetes_driver_container_name = kubernetes_driver_container_name def execute(self, context: Context) -> None: """Call the SparkSubmitHook to run the provided spark job.""" @@ -603,4 +607,5 @@ def _get_hook(self) -> SparkSubmitHook: track_driver_via_k8s_api=self._track_driver_via_k8s_api, yarn_track_via_rm_api=self._yarn_track_via_rm_api, yarn_rm_auth=self._yarn_rm_auth, + kubernetes_driver_container_name=self.kubernetes_driver_container_name, ) diff --git a/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_submit.py b/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_submit.py index d25ddbf7f76ee..57e9e9d5a27fa 100644 --- a/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_submit.py +++ b/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_submit.py @@ -22,12 +22,21 @@ from io import StringIO from pathlib import Path from types import ModuleType -from unittest.mock import MagicMock, call, mock_open, patch +from unittest.mock import ANY, MagicMock, call, mock_open, patch import kubernetes import pytest import requests -from kubernetes.client import V1Pod, V1PodStatus +from kubernetes.client import ( + V1Container, + V1ContainerState, + V1ContainerStateTerminated, + V1ContainerStateWaiting, + V1ContainerStatus, + V1Pod, + V1PodSpec, + V1PodStatus, +) from airflow.models import Connection from airflow.providers.apache.spark.hooks.spark_submit import SparkSubmitHook @@ -953,7 +962,7 @@ def test_process_spark_submit_log_k8s(self, pod_name): hook._process_spark_submit_log(log_lines) # Then - assert hook._kubernetes_driver_pod == pod_name + assert hook._kubernetes_driver_pod_name == pod_name assert hook._kubernetes_application_id == "spark-465b868ada474bda82ccb84ab2747fcd" assert hook._spark_exit_code == 999 @@ -978,7 +987,7 @@ def test_process_spark_submit_log_k8s_submission_id_format(self): hook._process_spark_submit_log(log_lines) - assert hook._kubernetes_driver_pod == "arrow-spark-c8e2e29e73db9c93-driver" + assert hook._kubernetes_driver_pod_name == "arrow-spark-c8e2e29e73db9c93-driver" def test_process_spark_client_mode_submit_log_k8s(self): # Given @@ -1250,7 +1259,7 @@ def test_on_kill_deletes_pod_when_k8s_api_tracking_and_submit_sp_already_exited( has already exited. """ hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) - hook._kubernetes_driver_pod = "spark-app-abc-driver" + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" hook._kubernetes_application_id = "spark-abc" hook._submit_sp = MagicMock() # spark-submit already exited @@ -1623,12 +1632,12 @@ def test_conf_injection_adds_wait_app_completion(self): @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") def test_poll_k8s_driver_succeeds(self, mock_get_client): hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) - hook._kubernetes_driver_pod = "spark-app-abc-driver" + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" hook._kubernetes_application_id = "spark-abc" mock_client = mock_get_client.return_value - running_pod = V1Pod(status=V1PodStatus(phase="Running")) - succeeded_pod = V1Pod(status=V1PodStatus(phase="Succeeded")) + running_pod = V1Pod(spec=V1PodSpec(containers=[]), status=V1PodStatus(phase="Running")) + succeeded_pod = V1Pod(spec=V1PodSpec(containers=[]), status=V1PodStatus(phase="Succeeded")) mock_client.read_namespaced_pod.side_effect = [running_pod, succeeded_pod] with patch.object(hook, "_run_post_submit_commands"): @@ -1639,11 +1648,11 @@ def test_poll_k8s_driver_succeeds(self, mock_get_client): @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") def test_poll_k8s_driver_raises_on_failed(self, mock_get_client): hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) - hook._kubernetes_driver_pod = "spark-app-abc-driver" + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" hook._kubernetes_application_id = "spark-abc" mock_client = mock_get_client.return_value - failed_pod = V1Pod(status=V1PodStatus(phase="Failed")) + failed_pod = V1Pod(spec=V1PodSpec(containers=[]), status=V1PodStatus(phase="Failed")) mock_client.read_namespaced_pod.return_value = failed_pod with pytest.raises(RuntimeError, match="phase=Failed"): @@ -1652,11 +1661,13 @@ def test_poll_k8s_driver_raises_on_failed(self, mock_get_client): @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") def test_poll_k8s_driver_raises_after_consecutive_unknown(self, mock_get_client): hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) - hook._kubernetes_driver_pod = "spark-app-abc-driver" + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" hook._kubernetes_application_id = "spark-abc" mock_client = mock_get_client.return_value - mock_client.read_namespaced_pod.return_value = V1Pod(status=V1PodStatus(phase="Unknown")) + mock_client.read_namespaced_pod.return_value = V1Pod( + spec=V1PodSpec(containers=[]), status=V1PodStatus(phase="Unknown") + ) with patch("time.sleep"), pytest.raises(RuntimeError, match="Unknown phase"): hook._poll_k8s_driver_via_api() @@ -1668,12 +1679,12 @@ def test_poll_k8s_driver_raises_after_consecutive_unknown(self, mock_get_client) @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") def test_poll_k8s_driver_tolerates_transient_api_errors(self, mock_get_client, _): hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) - hook._kubernetes_driver_pod = "spark-app-abc-driver" + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" hook._kubernetes_application_id = "spark-abc" mock_client = mock_get_client.return_value api_error = kube_client.ApiException(status=500, reason="Internal Server Error") - succeeded_pod = V1Pod(status=V1PodStatus(phase="Succeeded")) + succeeded_pod = V1Pod(spec=V1PodSpec(containers=[]), status=V1PodStatus(phase="Succeeded")) mock_client.read_namespaced_pod.side_effect = [api_error, api_error, succeeded_pod] with patch.object(hook, "_run_post_submit_commands"): @@ -1685,11 +1696,13 @@ def test_poll_k8s_driver_tolerates_transient_api_errors(self, mock_get_client, _ def test_post_submit_commands_run_exactly_once_on_k8s_path(self, mock_get_client): """_run_post_submit_commands must fire exactly once: in _poll_k8s_driver_via_api finally.""" hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) - hook._kubernetes_driver_pod = "spark-app-abc-driver" + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" hook._kubernetes_application_id = "spark-abc" mock_client = mock_get_client.return_value - mock_client.read_namespaced_pod.return_value = V1Pod(status=V1PodStatus(phase="Succeeded")) + mock_client.read_namespaced_pod.return_value = V1Pod( + spec=V1PodSpec(containers=[]), status=V1PodStatus(phase="Succeeded") + ) with patch.object(hook, "_run_post_submit_commands") as mock_cmd: hook._poll_k8s_driver_via_api() @@ -1700,7 +1713,7 @@ def test_post_submit_commands_run_exactly_once_on_k8s_path(self, mock_get_client @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") def test_poll_k8s_driver_raises_after_consecutive_api_errors(self, mock_get_client, _): hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) - hook._kubernetes_driver_pod = "spark-app-abc-driver" + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" hook._kubernetes_application_id = "spark-abc" mock_client = mock_get_client.return_value @@ -1716,7 +1729,7 @@ def test_poll_k8s_driver_raises_after_consecutive_api_errors(self, mock_get_clie def test_poll_k8s_driver_exits_cleanly_on_404(self, mock_get_client): """404 from read_namespaced_pod means pod was deleted by on_kill — should return cleanly, not raise.""" hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) - hook._kubernetes_driver_pod = "spark-app-abc-driver" + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" hook._kubernetes_application_id = "spark-abc" mock_client = mock_get_client.return_value @@ -1726,6 +1739,285 @@ def test_poll_k8s_driver_exits_cleanly_on_404(self, mock_get_client): mock_client.delete_namespaced_pod.assert_not_called() + @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") + def test_poll_k8s_driver_container_exit_zero_succeeds(self, mock_get_client): + """Driver container exits cleanly with code 0""" + hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" + hook._kubernetes_application_id = "spark-abc" + + mock_client = mock_get_client.return_value + terminated = V1ContainerStateTerminated(exit_code=0) + state = V1ContainerState(terminated=terminated) + container_status = V1ContainerStatus( + name="spark-driver", + state=state, + ready=False, + restart_count=0, + image="spark:3", + image_id="", + ) + pod = V1Pod( + spec=V1PodSpec(containers=[V1Container(name="spark-driver")]), + status=V1PodStatus(phase="Running", container_statuses=[container_status]), + ) + mock_client.read_namespaced_pod.return_value = pod + + with patch.object(hook, "_run_post_submit_commands"): + hook._poll_k8s_driver_via_api() + + assert mock_client.read_namespaced_pod.call_count == 1 + + @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") + def test_poll_k8s_driver_container_nonzero_exit_raises(self, mock_get_client): + """Driver container raises RuntimeError with non-zero exit code""" + hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" + hook._kubernetes_application_id = "spark-abc" + + mock_client = mock_get_client.return_value + terminated = V1ContainerStateTerminated(exit_code=1, reason="Error") + state = V1ContainerState(terminated=terminated) + container_status = V1ContainerStatus( + name="spark-driver", + state=state, + ready=False, + restart_count=0, + image="spark:3", + image_id="", + ) + pod = V1Pod( + spec=V1PodSpec(containers=[V1Container(name="spark-driver")]), + status=V1PodStatus(phase="Running", container_statuses=[container_status]), + ) + mock_client.read_namespaced_pod.return_value = pod + + with pytest.raises(RuntimeError, match="Exit code: 1"): + hook._poll_k8s_driver_via_api() + + @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") + def test_poll_k8s_driver_single_container_fallback(self, mock_get_client): + """Single container with no 'spark' or 'driver' in its name is still set as the driver container""" + hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" + hook._kubernetes_application_id = "spark-abc" + + mock_client = mock_get_client.return_value + terminated = V1ContainerStateTerminated(exit_code=0) + state = V1ContainerState(terminated=terminated) + container_status = V1ContainerStatus( + name="main", + state=state, + ready=False, + restart_count=0, + image="spark:3", + image_id="", + ) + pod = V1Pod( + spec=V1PodSpec(containers=[V1Container(name="main")]), + status=V1PodStatus(phase="Running", container_statuses=[container_status]), + ) + mock_client.read_namespaced_pod.return_value = pod + + with patch.object(hook, "_run_post_submit_commands"): + hook._poll_k8s_driver_via_api() + + assert mock_client.read_namespaced_pod.call_count == 1 + + @patch("time.sleep") + @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") + def test_poll_k8s_driver_prioritizes_driver_over_spark_sidecar(self, mock_get_client, _): + """A sidecar matching 'spark' that exits first must not be mistaken for the driver.""" + hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" + hook._kubernetes_application_id = "spark-abc" + + mock_client = mock_get_client.return_value + containers = [V1Container(name="spark-metrics-sidecar"), V1Container(name="app-driver")] + sidecar_done = V1ContainerStatus( + name="spark-metrics-sidecar", + state=V1ContainerState(terminated=V1ContainerStateTerminated(exit_code=0)), + ready=False, + restart_count=0, + image="busybox", + image_id="", + ) + driver_still_running = V1Pod( + spec=V1PodSpec(containers=containers), + status=V1PodStatus(phase="Running", container_statuses=[sidecar_done]), + ) + driver_done = V1ContainerStatus( + name="app-driver", + state=V1ContainerState(terminated=V1ContainerStateTerminated(exit_code=0)), + ready=False, + restart_count=0, + image="spark:3", + image_id="", + ) + driver_finished = V1Pod( + spec=V1PodSpec(containers=containers), + status=V1PodStatus(phase="Running", container_statuses=[sidecar_done, driver_done]), + ) + mock_client.read_namespaced_pod.side_effect = [driver_still_running, driver_finished] + + with patch.object(hook, "_run_post_submit_commands"), patch.object(hook.log, "info") as mock_info: + result = hook._poll_k8s_driver_via_api() + + assert result == "Succeeded" + assert mock_client.read_namespaced_pod.call_count == 2 + mock_info.assert_any_call("%s has been identified as the driver container", "app-driver") + + @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") + def test_poll_k8s_driver_custom_container_name_override(self, mock_get_client): + """An explicit kubernetes_driver_container_name matches by exact name, bypassing the heuristic.""" + hook = SparkSubmitHook( + conn_id="spark_k8s_cluster", + track_driver_via_k8s_api=True, + kubernetes_driver_container_name="custom-driver", + ) + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" + hook._kubernetes_application_id = "spark-abc" + + mock_client = mock_get_client.return_value + terminated_status = V1ContainerStatus( + name="custom-driver", + state=V1ContainerState(terminated=V1ContainerStateTerminated(exit_code=0)), + ready=False, + restart_count=0, + image="spark:3", + image_id="", + ) + pod = V1Pod( + spec=V1PodSpec( + containers=[V1Container(name="custom-driver"), V1Container(name="spark-exporter")] + ), + status=V1PodStatus(phase="Running", container_statuses=[terminated_status]), + ) + mock_client.read_namespaced_pod.return_value = pod + + with patch.object(hook, "_run_post_submit_commands"), patch.object(hook.log, "info") as mock_info: + result = hook._poll_k8s_driver_via_api() + + assert result == "Succeeded" + mock_info.assert_any_call("%s has been identified as the driver container", "custom-driver") + + @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") + def test_poll_k8s_driver_custom_container_name_no_match_raises(self, mock_get_client): + """An override that matches no container on the pod must raise, not fall back silently, + and must delete the now-orphaned driver pod before raising. + """ + hook = SparkSubmitHook( + conn_id="spark_k8s_cluster", + track_driver_via_k8s_api=True, + kubernetes_driver_container_name="does-not-exist", + ) + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" + hook._kubernetes_application_id = "spark-abc" + + mock_client = mock_get_client.return_value + pod = V1Pod( + spec=V1PodSpec(containers=[V1Container(name="spark-driver")]), + status=V1PodStatus(phase="Running"), + ) + mock_client.read_namespaced_pod.return_value = pod + + with pytest.raises(ValueError, match="does not match any of the containers in pod"): + hook._poll_k8s_driver_via_api() + + import kubernetes + + mock_client.delete_namespaced_pod.assert_called_once_with( + "spark-app-abc-driver", + "mynamespace", + body=kubernetes.client.V1DeleteOptions(), + pretty=True, + ) + + @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") + def test_poll_k8s_driver_failed_phase_with_completed_container_warns(self, mock_get_client): + """phase=Failed with a driver container that exited 0 must warn and succeed, not raise.""" + hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" + hook._kubernetes_application_id = "spark-abc" + + mock_client = mock_get_client.return_value + terminated_status = V1ContainerStatus( + name="spark-driver", + state=V1ContainerState(terminated=V1ContainerStateTerminated(exit_code=0)), + ready=False, + restart_count=0, + image="spark:3", + image_id="", + ) + pod = V1Pod( + spec=V1PodSpec(containers=[V1Container(name="spark-driver")]), + status=V1PodStatus(phase="Failed", container_statuses=[terminated_status]), + ) + mock_client.read_namespaced_pod.return_value = pod + + with ( + patch.object(hook, "_run_post_submit_commands"), + patch.object(hook.log, "warning") as mock_warning, + ): + result = hook._poll_k8s_driver_via_api() + + assert result == "Succeeded" + mock_warning.assert_any_call( + "Driver pod %s reported phase=Failed, but driver container %s exited 0; " + "treating the Spark application as succeeded.", + "spark-app-abc-driver", + "spark-driver", + ) + + @patch("time.sleep") + @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") + def test_poll_k8s_driver_container_waiting_warning(self, mock_get_client, _): + hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) + hook._kubernetes_driver_pod_name = "spark-app-abc-driver" + hook._kubernetes_application_id = "spark-abc" + + mock_client = mock_get_client.return_value + waiting_state = V1ContainerState(waiting=V1ContainerStateWaiting(reason="ContainerCreating")) + waiting_status = V1ContainerStatus( + name="spark-driver", + state=waiting_state, + ready=False, + restart_count=0, + image="spark:3", + image_id="", + ) + waiting_pod = V1Pod( + spec=V1PodSpec(containers=[V1Container(name="spark-driver")]), + status=V1PodStatus(phase="Running", container_statuses=[waiting_status]), + ) + terminated = V1ContainerStateTerminated(exit_code=0) + done_state = V1ContainerState(terminated=terminated) + done_status = V1ContainerStatus( + name="spark-driver", + state=done_state, + ready=False, + restart_count=0, + image="spark:3", + image_id="", + ) + done_pod = V1Pod( + spec=V1PodSpec(containers=[V1Container(name="spark-driver")]), + status=V1PodStatus(phase="Running", container_statuses=[done_status]), + ) + mock_client.read_namespaced_pod.side_effect = [waiting_pod] * 10 + [done_pod] + + with patch.object(hook, "_run_post_submit_commands"): + with patch.object(hook.log, "warning") as mock_warning: + hook._poll_k8s_driver_via_api() + + mock_warning.assert_any_call( + "Driver container %s has been waiting for %d polls (~%ds); " + "it may be unschedulable. Continuing to wait — set execution_timeout to bound wait time.", + "spark-driver", + 10, + ANY, + ) + @patch("airflow.providers.apache.spark.hooks.spark_submit.subprocess.run") def test_run_post_submit_commands_runs_only_once(self, mock_run): """Calling _run_post_submit_commands twice must execute commands exactly once.""" diff --git a/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py b/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py index c55cb10f3ebdd..c98e53adfcdb3 100644 --- a/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py +++ b/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py @@ -927,6 +927,16 @@ def _make_k8s_hook(self): hook._conf = {} return hook + def test_get_hook_passes_kubernetes_driver_container_name(self): + operator = self._make_operator(kubernetes_driver_container_name="custom-driver") + + hook = operator._get_hook() + + assert hook.kubernetes_driver_container_name == "custom-driver" + + def test_kubernetes_driver_container_name_is_templatable(self): + assert "kubernetes_driver_container_name" in SparkSubmitOperator.template_fields + def test_execute_calls_submit_then_poll_when_flag_set(self): operator = self._make_operator(track_driver_via_k8s_api=True) hook = self._make_k8s_hook() @@ -956,7 +966,7 @@ def test_execute_falls_through_to_plain_submit_when_flag_off(self): def test_k8s_submit_job_returns_encoded_external_id(self): operator = self._make_operator(track_driver_via_k8s_api=True) hook = self._make_k8s_hook() - hook._kubernetes_driver_pod = "spark-abc-driver" + hook._kubernetes_driver_pod_name = "spark-abc-driver" hook._connection = {"namespace": "mynamespace"} operator._hook = hook @@ -969,7 +979,7 @@ def test_k8s_submit_job_returns_encoded_external_id(self): def test_k8s_submit_job_raises_when_pod_name_missing(self): operator = self._make_operator(track_driver_via_k8s_api=True) hook = self._make_k8s_hook() - hook._kubernetes_driver_pod = None + hook._kubernetes_driver_pod_name = None hook._connection = {"namespace": "mynamespace"} operator._hook = hook @@ -1050,7 +1060,7 @@ def test_k8s_poll_until_complete_sets_pod_name_and_calls_poll_api(self): operator.poll_until_complete("mynamespace:spark-abc-driver", {}) - assert hook._kubernetes_driver_pod == "spark-abc-driver" + assert hook._kubernetes_driver_pod_name == "spark-abc-driver" hook._poll_k8s_driver_via_api.assert_called_once() @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="task_state_store requires Airflow 3.3+") @@ -1103,7 +1113,7 @@ def test_k8s_execute_persists_pod_id_when_durable(self): """execute() with durable=True stores the pod ID in task_store before polling.""" operator = self._make_operator(track_driver_via_k8s_api=True, durable=True) hook = self._make_k8s_hook() - hook._kubernetes_driver_pod = "spark-abc-driver" + hook._kubernetes_driver_pod_name = "spark-abc-driver" hook._connection = {"namespace": "mynamespace"} operator._hook = hook task_store = FakeTaskStateStore() @@ -1126,7 +1136,7 @@ def test_k8s_execute_durable_false_does_not_persist_pod_id(self): """execute() with durable=False does not write spark_job_id to task_store.""" operator = self._make_operator(track_driver_via_k8s_api=True, durable=False) hook = self._make_k8s_hook() - hook._kubernetes_driver_pod = "spark-abc-driver" + hook._kubernetes_driver_pod_name = "spark-abc-driver" hook._connection = {"namespace": "mynamespace"} operator._hook = hook task_store = FakeTaskStateStore()