From b2f919075cab0ff80fc6d21b4584e9636a8df3be Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Tue, 21 Jul 2026 16:53:04 +0530 Subject: [PATCH 1/2] Surface YARN diagnostics for spark on yarn clusters --- .../apache/spark/hooks/spark_submit.py | 21 ++++++--- .../apache/spark/hooks/test_spark_submit.py | 45 ++++++++++++++++++- 2 files changed, 58 insertions(+), 8 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 662966e4e1459..7b6200e374dec 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 @@ def _start_yarn_application_status_tracking(self, application_id: str) -> None: 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,16 +933,18 @@ def _start_yarn_application_status_tracking(self, application_id: str) -> None: 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( @@ -1001,8 +1003,15 @@ def _resolved_yarn_rm_auth(self) -> AuthBase | None: 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 +1026,7 @@ def _query_yarn_application_status(self, application_id: str) -> tuple[str, str] ) 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 " @@ -1360,7 +1369,7 @@ def query_yarn_application_status(self, application_id: str) -> str: - 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 90f923905fa7b..839395fca4069 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 @@ def _rm_kill_url(cls, app_id: str | None = None) -> str: 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,42 @@ def test_yarn_status_tracking_fails_on_failed_state_with_undefined_final_status( 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", 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): From 34c1c26d49eebb3c64cf5d2a1c556997c294b03e Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Wed, 22 Jul 2026 15:57:51 +0530 Subject: [PATCH 2/2] comments from kaxil --- .../providers/apache/spark/hooks/spark_submit.py | 3 ++- .../unit/apache/spark/hooks/test_spark_submit.py | 14 +++++++++++++- 2 files changed, 15 insertions(+), 2 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 7b6200e374dec..924cb8432a280 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 @@ -948,7 +948,8 @@ def _start_yarn_application_status_tracking(self, application_id: str) -> None: ) 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) 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 839395fca4069..5733f44fd6b27 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 @@ -1783,7 +1783,9 @@ def test_yarn_status_tracking_fails_on_failed_state_with_undefined_final_status( 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", diagnostics="Application application_1700000000000_0001 was killed by user root" + "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) @@ -1826,6 +1828,16 @@ def test_yarn_status_tracking_fails_on_unexpected_final_status(self, mock_get, m 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()."""