kaxil commented on code in PR #67715:
URL: https://github.com/apache/airflow/pull/67715#discussion_r3336901894
##########
providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py:
##########
@@ -802,6 +861,95 @@ def _start_driver_status_tracking(self) -> None:
f"returncode = {returncode}"
)
+ def _poll_k8s_driver_via_api(self) -> None:
+ """Poll the K8s driver pod phase until it reaches a terminal state."""
+ pod_name = self._kubernetes_driver_pod
+ namespace = self._connection["namespace"]
+ app_id = self._kubernetes_application_id or pod_name
+
+ if not pod_name:
+ raise ValueError("K8s driver pod name not set; cannot poll
status.")
+
+ client = kube_client.get_kube_client()
+ poll_interval = max(self._status_poll_interval, 20)
+ if poll_interval != self._status_poll_interval:
+ self.log.info(
+ "status_poll_interval=%ds is below the 20s minimum for K8s API
polling; using 20s.",
+ self._status_poll_interval,
+ )
+ # Mirror `missed_job_status_reports` / `max_missed_job_status_reports`
from
+ # `_start_driver_status_tracking`: tolerate transient failures before
giving up.
+ consecutive_unknown = 0
+ max_consecutive_unknown = 3
+ consecutive_api_errors = 0
+ max_consecutive_api_errors = 3
+ consecutive_pending = 0
+ pending_warn_threshold = 10
+
+ try:
+ while True:
+ try:
+ pod = client.read_namespaced_pod(pod_name, namespace)
+ consecutive_api_errors = 0
+ except kube_client.ApiException as e:
+ consecutive_api_errors += 1
+ self.log.warning(
+ "ApiException polling pod %s (%d/%d): %s",
+ pod_name,
+ consecutive_api_errors,
+ max_consecutive_api_errors,
+ e,
+ )
+ if consecutive_api_errors >= max_consecutive_api_errors:
+ raise RuntimeError(
+ f"K8s API unreachable after
{consecutive_api_errors} consecutive errors "
+ f"while polling {app_id}; giving up."
+ ) from e
+ 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,
+ )
+ 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)
+ try:
+ client.delete_namespaced_pod(pod_name, namespace)
+ self.log.info("Deleted driver pod %s", pod_name)
+ except kube_client.ApiException:
+ self.log.warning("Could not delete driver pod %s after
completion", pod_name)
+ finally:
+ self._run_post_submit_commands()
Review Comment:
This is the second `_run_post_submit_commands()` call on the K8s-API path.
The pre-existing `finally` in `submit()` (hooks/spark_submit.py:706) runs them
whenever `not self._should_track_driver_status`, and
`_should_track_driver_status` is False for K8s tracking (it's only set for
`spark://` cluster mode). So `submit()` already fires the commands right after
`spark-submit` exits, which under `waitAppCompletion=false` is seconds after
pod creation while the driver is still running; this `finally` then fires them
again at terminal state. For the documented Istio `quitquitquit` case the first
call tears the sidecar down mid-poll. Suggest gating the `submit()` finally on
`_should_track_driver_via_k8s_api()` too, so the commands run only once (here).
(`test_poll_k8s_driver_succeeds` patches `_run_post_submit_commands` away, so
the suite can't catch the double-run; a test asserting it runs exactly once
across the full K8s path would lock the fix.)
##########
providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py:
##########
@@ -227,13 +235,21 @@ def execute(self, context: Context) -> None:
if self._hook is None:
self._hook = self._get_hook()
hook = self._hook
+ if self._track_driver_via_k8s_api:
+ hook._validate_track_driver_via_k8s_api_config()
if hook._should_track_driver_status:
if self.reconnect_on_retry:
return self.execute_resumable(context)
# reconnect_on_retry=False: still submit-and-poll, just skip
task_state persistence.
driver_id = self.submit_job(context)
self.poll_until_complete(driver_id, context)
return self.get_job_result(driver_id, context)
+ if hook._should_track_driver_via_k8s_api():
+ # TODO: Wire into execute_resumable() via ResumableJobMixin
+ # (fill submit_job / poll_until_complete K8s stubs) to enable
crash recovery.
+ hook.submit(self.application)
+ hook._poll_k8s_driver_via_api()
Review Comment:
A task kill during this poll won't clean up the driver pod. `on_kill`'s K8s
pod-deletion block (hooks/spark_submit.py:986) is nested under `if
self._submit_sp and self._submit_sp.poll() is None:`. By the time
`_poll_k8s_driver_via_api` is running, `submit()` has already `.wait()`ed on
`spark-submit`, which exits early under `waitAppCompletion=false`, so `poll()`
is no longer `None` and the deletion block is skipped. A kill in the polling
window (timeout, manual clear, worker SIGTERM) therefore orphans the running
Spark driver. Consider deleting the pod in `on_kill` whenever
`_should_track_driver_via_k8s_api()` is true and `_kubernetes_driver_pod` is
set, independent of `_submit_sp` state.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
