amoghrajesh commented on code in PR #67715:
URL: https://github.com/apache/airflow/pull/67715#discussion_r3338792204
##########
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:
Handled in d12200d5f3
##########
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:
Handled in d12200d5f3
##########
providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py:
##########
@@ -268,6 +282,34 @@ def _resolve_should_track_driver_status(self) -> bool:
"""
return "spark://" in self._connection["master"] and
self._connection["deploy_mode"] == "cluster"
+ def _should_track_driver_via_k8s_api(self) -> bool:
+ return (
+ self._track_driver_via_k8s_api
+ and self._is_kubernetes
+ and self._connection["deploy_mode"] == "cluster"
+ )
+
+ def _validate_track_driver_via_k8s_api_config(self) -> None:
+ if not self._is_kubernetes:
+ raise ValueError(
+ "`track_driver_via_k8s_api=True` requires Spark master to be
Kubernetes (k8s://...)."
+ )
+ if self._connection["deploy_mode"] != "cluster":
+ raise ValueError(
+ "`track_driver_via_k8s_api=True` requires
`deploy_mode='cluster'`; "
+ f"got deploy_mode={self._connection['deploy_mode']!r}."
+ )
+ if not self._connection.get("namespace"):
+ raise ValueError(
+ "`track_driver_via_k8s_api=True` requires a namespace; "
+ "set it in the connection extra as `namespace` or via
`spark.kubernetes.namespace` in conf."
+ )
+ if self._conf.get(_K8S_WAIT_APP_COMPLETION_CONF, "").lower() == "true":
Review Comment:
Handled in d12200d5f3
##########
providers/apache/spark/docs/operators.rst:
##########
@@ -214,3 +214,32 @@ See :doc:`connections/spark-submit` for how to configure
these fields.
.. note::
Crash recovery in cluster mode requires Airflow 3.3+ (``task_state``
support). On earlier
versions the operator falls back to the previous behavior of always
submitting fresh.
+
+Tracking driver status via Kubernetes API
+""""""""""""""""""""""""""""""""""""""""""
+
+When running in Kubernetes cluster mode, ``spark-submit`` blocks for the
duration of the job.
+The JVM runs processes which does nothing but polling of the pod phase and
holds heap space for
+the entire duration. This is not ideal for long-running jobs, especially when
the driver is idle
+for long periods (e.g. waiting for data or user input).
+
+Set ``track_driver_via_k8s_api=True`` to have the operator track the driver
pod status via the
+Python Kubernetes client rather than holding ``spark-submit`` open for the
full job duration:
+
+.. code-block:: python
+
+ from airflow.providers.apache.spark.operators.spark_submit import
SparkSubmitOperator
+
+ 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,
+ )
+
+**Requirements**
+
+* The Spark connection ``master`` must be ``k8s://...`` and ``deploy_mode``
must be ``cluster``.
+* Do not set ``spark.kubernetes.submission.waitAppCompletion=true`` in your
``conf`` — this
+ conflicts with the flag and a ``ValueError`` will be raised at task start.
Review Comment:
Handled in d12200d5f3
--
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]
