Surface YARN diagnostics for spark on yarn clusters

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,20 +933,23 @@ 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(
-                    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 @@ 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 +1027,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 +1370,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:

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,44 @@ 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",
+            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 @@ 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()."""