From a4aab1c778148a269c826f80361376641f856770 Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Thu, 4 Jun 2026 20:48:25 -0400 Subject: [PATCH 1/7] Use container state to determine spark job status --- .../apache/spark/hooks/spark_submit.py | 94 +++++++++++++------ 1 file changed, 63 insertions(+), 31 deletions(-) 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 7cf1f3248adc8..f419927c56c73 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 @@ -1130,7 +1130,7 @@ def _poll_k8s_driver_via_api(self) -> None: consecutive_api_errors = 0 max_consecutive_api_errors = 3 consecutive_pending = 0 - pending_warn_threshold = 10 + waiting_or_pending_warn_threshold = 10 try: if not pod_name: @@ -1162,39 +1162,71 @@ def _poll_k8s_driver_via_api(self) -> None: time.sleep(poll_interval) continue - phase = pod.status.phase or "Initializing" - self.log.info("Application status for %s (phase: %s)", app_id, phase) - if phase == "Succeeded": - break - if phase == "Failed": - 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 == "Pending": - consecutive_pending += 1 - if consecutive_pending == 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, - consecutive_pending, - consecutive_pending * poll_interval, - ) + for container in pod.spec.containers: + if "spark" in container.name.lower() or "driver" in container.name.lower(): + driver_container = container + break + if len(pod.spec.containers) == 1: + driver_container = container + else: + driver_container + None + + if driver_container: + 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: + 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_pending = 0 - - if phase == "Unknown": - consecutive_unknown += 1 - if consecutive_unknown >= max_consecutive_unknown: + phase = pod.status.phase or "Initializing" + self.log.info("Application status for %s (phase: %s)", app_id, phase) + if phase == "Succeeded": + break + if phase == "Failed": + 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} reported Unknown phase " - f"{consecutive_unknown} times consecutively; giving up." + f"Spark application {app_id} failed (phase=Failed{container_state})" ) - else: - consecutive_unknown = 0 + if phase == "Pending": + consecutive_pending += 1 + 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, + consecutive_pending, + consecutive_pending * poll_interval, + ) + else: + consecutive_pending = 0 + + if phase == "Unknown": + consecutive_unknown += 1 + if consecutive_unknown >= max_consecutive_unknown: + raise RuntimeError( + f"Spark application {app_id} reported Unknown phase " + f"{consecutive_unknown} times consecutively; giving up." + ) + else: + consecutive_unknown = 0 time.sleep(poll_interval) self._delete_driver_pod() finally: From 6f48fa5e9564c8ecd16ad64205b5c96f401c0c50 Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Thu, 4 Jun 2026 20:54:52 -0400 Subject: [PATCH 2/7] Correct typo and conditional logic --- .../src/airflow/providers/apache/spark/hooks/spark_submit.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 f419927c56c73..f427a2f0bef8e 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 @@ -1169,7 +1169,7 @@ def _poll_k8s_driver_via_api(self) -> None: if len(pod.spec.containers) == 1: driver_container = container else: - driver_container + None + driver_container = None if driver_container: for status in pod.status.container_statuses or []: @@ -1191,6 +1191,8 @@ def _poll_k8s_driver_via_api(self) -> None: consecutive_waiting, consecutive_waiting * poll_interval, ) + else: + consecutive_waiting = 0 else: phase = pod.status.phase or "Initializing" self.log.info("Application status for %s (phase: %s)", app_id, phase) From 5626137b9fe07bc70753ce25ee845ee34941bdb0 Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Tue, 9 Jun 2026 11:01:21 -0400 Subject: [PATCH 3/7] Correct failing tests --- .../apache/spark/hooks/spark_submit.py | 4 ++-- .../apache/spark/hooks/test_spark_submit.py | 18 +++++++++++------- 2 files changed, 13 insertions(+), 9 deletions(-) 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 f427a2f0bef8e..f6e767d4d45e0 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 @@ -1130,6 +1130,7 @@ def _poll_k8s_driver_via_api(self) -> None: consecutive_api_errors = 0 max_consecutive_api_errors = 3 consecutive_pending = 0 + consecutive_waiting = 0 waiting_or_pending_warn_threshold = 10 try: @@ -1161,6 +1162,7 @@ def _poll_k8s_driver_via_api(self) -> None: ) from e time.sleep(poll_interval) continue + driver_container = None for container in pod.spec.containers: if "spark" in container.name.lower() or "driver" in container.name.lower(): @@ -1168,8 +1170,6 @@ def _poll_k8s_driver_via_api(self) -> None: break if len(pod.spec.containers) == 1: driver_container = container - else: - driver_container = None if driver_container: for status in pod.status.container_statuses or []: 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 f4a610a940814..23fea8ce285de 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 @@ -27,7 +27,7 @@ import kubernetes import pytest import requests -from kubernetes.client import V1Pod, V1PodStatus +from kubernetes.client import V1Pod, V1PodSpec, V1PodStatus from airflow.models import Connection from airflow.providers.apache.spark.hooks.spark_submit import SparkSubmitHook @@ -1503,8 +1503,8 @@ def test_poll_k8s_driver_succeeds(self, mock_get_client): 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"): @@ -1519,7 +1519,7 @@ def test_poll_k8s_driver_raises_on_failed(self, mock_get_client): 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"): @@ -1532,7 +1532,9 @@ def test_poll_k8s_driver_raises_after_consecutive_unknown(self, mock_get_client) 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() @@ -1549,7 +1551,7 @@ def test_poll_k8s_driver_tolerates_transient_api_errors(self, mock_get_client, _ 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"): @@ -1565,7 +1567,9 @@ def test_post_submit_commands_run_exactly_once_on_k8s_path(self, mock_get_client 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() From 5b713894cbcf35acbbb27018731ead682b6fd70b Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Tue, 9 Jun 2026 12:32:32 -0400 Subject: [PATCH 4/7] Add unit tests and fix bugs --- .../apache/spark/hooks/spark_submit.py | 5 +- .../apache/spark/hooks/test_spark_submit.py | 147 +++++++++++++++++- 2 files changed, 149 insertions(+), 3 deletions(-) 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 f6e767d4d45e0..1d08117758a67 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 @@ -1170,13 +1170,14 @@ def _poll_k8s_driver_via_api(self) -> None: break if len(pod.spec.containers) == 1: driver_container = container - + container_completed = False if driver_container: 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}" @@ -1229,6 +1230,8 @@ def _poll_k8s_driver_via_api(self) -> None: ) else: consecutive_unknown = 0 + if container_completed: + break time.sleep(poll_interval) self._delete_driver_pod() finally: 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 23fea8ce285de..112b751b5dc6d 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, V1PodSpec, 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 @@ -1606,6 +1615,140 @@ 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 = "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 = "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 = "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_container_waiting_warning(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_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.""" From 192d612f0f4f297568d7855a61ee38f2a422c597 Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Wed, 17 Jun 2026 00:26:52 -0400 Subject: [PATCH 5/7] Fix failing tests --- .../airflow/providers/apache/spark/hooks/spark_submit.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 2fcd882be19ad..dbca346d71d83 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 @@ -1166,6 +1166,7 @@ def _poll_k8s_driver_via_api(self) -> str | None: consecutive_pending = 0 consecutive_waiting = 0 waiting_or_pending_warn_threshold = 10 + terminal_phase: str | None = None try: if not pod_name: @@ -1254,7 +1255,7 @@ def _poll_k8s_driver_via_api(self) -> str | None: raise RuntimeError(f"Spark application {app_id} failed (phase=Failed{container_state})") 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.", @@ -1299,6 +1300,9 @@ 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 From ef669f35ad1742b9eba658b1ba49ff07974dc553 Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Tue, 7 Jul 2026 13:37:10 -0400 Subject: [PATCH 6/7] Remove redundant code --- .../apache/spark/hooks/spark_submit.py | 46 +++++-------------- 1 file changed, 12 insertions(+), 34 deletions(-) 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 dbca346d71d83..c98e8488de017 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 @@ -1229,44 +1229,22 @@ def _poll_k8s_driver_via_api(self) -> str | None: ) 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": - if pod.status.container_statuses: - cs = pod.status.container_statuses[0] - if cs.state and cs.state.terminated: - t = cs.state.terminated - self.log.info( - "Container final status: exit_code=%s reason=%s started_at=%s finished_at=%s", - t.exit_code, - t.reason, - t.started_at, - t.finished_at, - ) - terminal_phase = phase - break - if phase == "Failed": - 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 == "Pending": - consecutive_pending += 1 - 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, - consecutive_pending, - consecutive_pending * poll_interval, - ) else: phase = pod.status.phase or "Initializing" self.log.info("Application status for %s (phase: %s)", app_id, phase) if phase == "Succeeded": + if pod.status.container_statuses: + cs = pod.status.container_statuses[0] + if cs.state and cs.state.terminated: + t = cs.state.terminated + self.log.info( + "Container final status: exit_code=%s reason=%s started_at=%s finished_at=%s", + t.exit_code, + t.reason, + t.started_at, + t.finished_at, + ) + terminal_phase = phase break if phase == "Failed": container_state = "" From 329d105a34a45e8bc907d30b36ae01da42e76656 Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Mon, 20 Jul 2026 20:55:13 -0400 Subject: [PATCH 7/7] Safeguard against unschedulable pods --- .../apache/spark/hooks/spark_submit.py | 89 +++++++++---------- 1 file changed, 43 insertions(+), 46 deletions(-) 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 c98e8488de017..5dfd2d27c3eeb 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 @@ -1140,7 +1140,7 @@ 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``). @@ -1229,54 +1229,51 @@ def _poll_k8s_driver_via_api(self) -> str | None: ) 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": + if pod.status.container_statuses: + cs = pod.status.container_statuses[0] + if cs.state and cs.state.terminated: + t = cs.state.terminated + self.log.info( + "Container final status: exit_code=%s reason=%s started_at=%s finished_at=%s", + t.exit_code, + t.reason, + t.started_at, + t.finished_at, + ) + terminal_phase = phase + break + 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 == "Pending": + consecutive_pending += 1 + 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, + consecutive_pending, + consecutive_pending * poll_interval, + ) else: - phase = pod.status.phase or "Initializing" - self.log.info("Application status for %s (phase: %s)", app_id, phase) - if phase == "Succeeded": - if pod.status.container_statuses: - cs = pod.status.container_statuses[0] - if cs.state and cs.state.terminated: - t = cs.state.terminated - self.log.info( - "Container final status: exit_code=%s reason=%s started_at=%s finished_at=%s", - t.exit_code, - t.reason, - t.started_at, - t.finished_at, - ) - terminal_phase = phase - break - if phase == "Failed": - 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}" + consecutive_pending = 0 + + if phase == "Unknown": + consecutive_unknown += 1 + if consecutive_unknown >= max_consecutive_unknown: raise RuntimeError( - f"Spark application {app_id} failed (phase=Failed{container_state})" + f"Spark application {app_id} reported Unknown phase " + f"{consecutive_unknown} times consecutively; giving up." ) - if phase == "Pending": - consecutive_pending += 1 - 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, - consecutive_pending, - consecutive_pending * poll_interval, - ) - else: - consecutive_pending = 0 - - if phase == "Unknown": - consecutive_unknown += 1 - if consecutive_unknown >= max_consecutive_unknown: - raise RuntimeError( - f"Spark application {app_id} reported Unknown phase " - f"{consecutive_unknown} times consecutively; giving up." - ) - else: - consecutive_unknown = 0 + 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.