This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new c173545e8d0 Surface YARN diagnostics for spark on yarn clusters
(#70183)
c173545e8d0 is described below
commit c173545e8d0ffb45274855332ce39ab166698e7b
Author: Amogh Desai <[email protected]>
AuthorDate: Wed Jul 22 17:10:41 2026 +0530
Surface YARN diagnostics for spark on yarn clusters (#70183)
---
.../providers/apache/spark/hooks/spark_submit.py | 24 ++++++---
.../unit/apache/spark/hooks/test_spark_submit.py | 57 +++++++++++++++++++++-
2 files changed, 72 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 6d90f7b003b..cd4a0703d19 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
@@ -907,7 +907,7 @@ class SparkSubmitHook(BaseHook, LoggingMixin):
while True:
self.log.debug("Polling YARN RM REST API for application %s",
application_id)
try:
- state, final_status =
self._query_yarn_application_status(application_id)
+ state, final_status, diagnostics =
self._query_yarn_application_status(application_id)
except RuntimeError as exc:
consecutive_failures += 1
if consecutive_failures > max_consecutive_failures:
@@ -933,20 +933,23 @@ class SparkSubmitHook(BaseHook, LoggingMixin):
elif poll_count % heartbeat_interval == 0:
self.log.info("YARN application %s is still %s",
application_id, state)
+ diagnostics_suffix = f"\nDiagnostics: {diagnostics}" if
diagnostics else ""
if state in self._YARN_FINAL_FAILURES:
raise RuntimeError(
f"YARN application {application_id} ended with state:
{state}, "
- f"final status: {final_status}"
+ f"final status: {final_status}{diagnostics_suffix}"
)
if final_status == self._YARN_FINAL_SUCCESS:
return
if final_status in self._YARN_FINAL_FAILURES:
raise RuntimeError(
f"YARN application {application_id} ended with final
status: {final_status}"
+ f"{diagnostics_suffix}"
)
if final_status != self._YARN_FINAL_UNDEFINED:
raise RuntimeError(
- f"YARN application {application_id} returned unexpected
final status: {final_status}"
+ f"YARN application {application_id} returned unexpected
final status: "
+ f"{final_status}{diagnostics_suffix}"
)
time.sleep(poll_interval)
@@ -1001,8 +1004,15 @@ class SparkSubmitHook(BaseHook, LoggingMixin):
return None
- def _query_yarn_application_status(self, application_id: str) ->
tuple[str, str]:
- """GET ``/ws/v1/cluster/apps/{id}`` once and return ``app.state`` and
``app.finalStatus``."""
+ def _query_yarn_application_status(self, application_id: str) ->
tuple[str, str, str]:
+ """
+ GET ``/ws/v1/cluster/apps/{id}`` once.
+
+ Returns ``app.state``, ``app.finalStatus``, and ``app.diagnostics`` -
diagnostics is
+ where YARN puts the actual human readable failure reason (AM launch
error, container
+ OOM, explicit kill, etc.), so failure exceptions can include it
instead of just the
+ two terminal-state enum values.
+ """
url =
f"{self._get_yarn_rm_base_url()}/ws/v1/cluster/apps/{application_id}"
try:
resp = requests.get(url, auth=self._resolved_yarn_rm_auth,
timeout=self._HTTP_TIMEOUT)
@@ -1017,7 +1027,7 @@ class SparkSubmitHook(BaseHook, LoggingMixin):
)
try:
app = resp.json()["app"]
- return app["state"], app["finalStatus"]
+ return app["state"], app["finalStatus"], app.get("diagnostics", "")
except (ValueError, KeyError, TypeError) as exc:
raise RuntimeError(
f"YARN RM REST API returned unexpected payload for application
"
@@ -1362,7 +1372,7 @@ class SparkSubmitHook(BaseHook, LoggingMixin):
- FINISHED + any other finalStatus -> "FAILED"
- FAILED or KILLED -> "FAILED"
"""
- state, final_status =
self._query_yarn_application_status(application_id)
+ state, final_status, _ =
self._query_yarn_application_status(application_id)
if state in {"NEW", "NEW_SAVING", "SUBMITTED", "ACCEPTED", "RUNNING"}:
return state
if state == "FINISHED" and final_status == self._YARN_FINAL_SUCCESS:
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 90f923905fa..5733f44fd6b 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
@@ -1666,10 +1666,15 @@ class TestSparkSubmitHook:
return f"{cls._RM_BASE_URL}/ws/v1/cluster/apps/{app_id or
cls._RM_APP_ID}/state"
@classmethod
- def _rm_status_resp(cls, final_status: str, state: str = "FINISHED") ->
MagicMock:
+ def _rm_status_resp(
+ cls, final_status: str, state: str = "FINISHED", diagnostics: str |
None = None
+ ) -> MagicMock:
resp = MagicMock(spec=requests.Response)
resp.status_code = 200
- resp.json.return_value = {"app": {"id": cls._RM_APP_ID, "state":
state, "finalStatus": final_status}}
+ app = {"id": cls._RM_APP_ID, "state": state, "finalStatus":
final_status}
+ if diagnostics is not None:
+ app["diagnostics"] = diagnostics
+ resp.json.return_value = {"app": app}
return resp
@staticmethod
@@ -1773,6 +1778,44 @@ class TestSparkSubmitHook:
mock_sleep.assert_not_called()
+ @patch("airflow.providers.apache.spark.hooks.spark_submit.time.sleep")
+ @patch("airflow.providers.apache.spark.hooks.spark_submit.requests.get")
+ def test_yarn_status_tracking_includes_diagnostics_on_state_failure(self,
mock_get, mock_sleep):
+ """RM state FAILED/KILLED -> raised message includes the RM's
diagnostics field."""
+ mock_get.return_value = self._rm_status_resp(
+ "KILLED",
+ state="KILLED",
+ diagnostics="Application application_1700000000000_0001 was killed
by user root",
+ )
+
+ hook = SparkSubmitHook(conn_id="spark_yarn_rm",
yarn_track_via_rm_api=True)
+ with pytest.raises(RuntimeError, match="Diagnostics: Application .*
was killed by user root"):
+ hook._start_yarn_application_status_tracking(self._RM_APP_ID)
+
+ @patch("airflow.providers.apache.spark.hooks.spark_submit.time.sleep")
+ @patch("airflow.providers.apache.spark.hooks.spark_submit.requests.get")
+ def
test_yarn_status_tracking_includes_diagnostics_on_final_status_failure(self,
mock_get, mock_sleep):
+ """RM finalStatus FAILED (state FINISHED) -> raised message includes
diagnostics."""
+ mock_get.return_value = self._rm_status_resp(
+ "FAILED", diagnostics="AM Container exited with exitCode: 1"
+ )
+
+ hook = SparkSubmitHook(conn_id="spark_yarn_rm",
yarn_track_via_rm_api=True)
+ with pytest.raises(RuntimeError, match="Diagnostics: AM Container
exited with exitCode: 1"):
+ hook._start_yarn_application_status_tracking(self._RM_APP_ID)
+
+ @patch("airflow.providers.apache.spark.hooks.spark_submit.time.sleep")
+ @patch("airflow.providers.apache.spark.hooks.spark_submit.requests.get")
+ def test_yarn_status_tracking_omits_diagnostics_suffix_when_absent(self,
mock_get, mock_sleep):
+ """RM response with no diagnostics field -> message has no
'Diagnostics:' suffix."""
+ mock_get.return_value = self._rm_status_resp("KILLED")
+
+ hook = SparkSubmitHook(conn_id="spark_yarn_rm",
yarn_track_via_rm_api=True)
+ with pytest.raises(RuntimeError) as exc_info:
+ hook._start_yarn_application_status_tracking(self._RM_APP_ID)
+
+ assert "Diagnostics:" not in str(exc_info.value)
+
@patch("airflow.providers.apache.spark.hooks.spark_submit.time.sleep")
@patch("airflow.providers.apache.spark.hooks.spark_submit.requests.get")
def test_yarn_status_tracking_fails_on_unexpected_final_status(self,
mock_get, mock_sleep):
@@ -1785,6 +1828,16 @@ class TestSparkSubmitHook:
mock_sleep.assert_not_called()
+ @patch("airflow.providers.apache.spark.hooks.spark_submit.time.sleep")
+ @patch("airflow.providers.apache.spark.hooks.spark_submit.requests.get")
+ def
test_yarn_status_tracking_includes_diagnostics_on_unexpected_final_status(self,
mock_get, mock_sleep):
+ """RM returns a non-standard finalStatus -> raised message also
includes diagnostics."""
+ mock_get.return_value = self._rm_status_resp("ENDED",
diagnostics="Application state is ENDED")
+
+ hook = SparkSubmitHook(conn_id="spark_yarn_rm",
yarn_track_via_rm_api=True)
+ with pytest.raises(RuntimeError, match="unexpected final status:
ENDED\nDiagnostics: .*ENDED"):
+ hook._start_yarn_application_status_tracking(self._RM_APP_ID)
+
@patch("airflow.providers.apache.spark.hooks.spark_submit.subprocess.Popen")
def
test_yarn_submit_captures_app_id_without_submitted_application_log(self,
mock_popen):
"""App ID parsed from log lines other than 'Submitted application ...'
is captured by hook.submit()."""
