feat: add yarn client (status, logs, kill)

spark_executor/core/yarn_client.py

@@ -0,0 +1,47 @@
+# coding=utf-8
+"""
+@Time   :2026/6/24
+@Author :tao.chen
+"""
+import re
+import subprocess
+
+
+class YarnError(Exception):
+    """Raised when a yarn CLI invocation fails."""
+
+
+_STATE_RE = re.compile(r"State\s*:\s*(\S+)")
+
+
+def _run(cmd: list[str]) -> "subprocess.CompletedProcess[str]":
+    return subprocess.run(cmd, capture_output=True, text=True, errors="replace")
+
+
+def get_application_status(application_id: str) -> tuple[str, str]:
+    proc = _run(["yarn", "application", "-status", application_id])
+    if proc.returncode != 0:
+        raise YarnError(
+            f"yarn application -status failed (rc={proc.returncode}): {proc.stderr}"
+        )
+    match = _STATE_RE.search(proc.stdout)
+    if not match:
+        raise YarnError(f"Could not parse YARN state from output: {proc.stdout!r}")
+    return match.group(1), proc.stdout
+
+
+def get_application_logs(application_id: str) -> str:
+    proc = _run(["yarn", "logs", "-applicationId", application_id])
+    if proc.returncode != 0:
+        raise YarnError(
+            f"yarn logs failed (rc={proc.returncode}): {proc.stderr}"
+        )
+    return proc.stdout
+
+
+def kill_application(application_id: str) -> None:
+    proc = _run(["yarn", "application", "-kill", application_id])
+    if proc.returncode != 0:
+        raise YarnError(
+            f"yarn application -kill failed (rc={proc.returncode}): {proc.stderr}"
+        )

tests/unit/test_yarn_client.py

@@ -0,0 +1,63 @@
+# coding=utf-8
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from spark_executor.core.yarn_client import (
+    YarnError,
+    get_application_logs,
+    get_application_status,
+    kill_application,
+)
+
+
+def _fake_proc(returncode: int, stdout: str = "", stderr: str = ""):
+    p = MagicMock()
+    p.returncode = returncode
+    p.stdout = stdout
+    p.stderr = stderr
+    return p
+
+
+def test_status_parses_state_line():
+    fake = _fake_proc(
+        0,
+        stdout="Application Report :\n  State : RUNNING\n  ...\n",
+    )
+    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
+        state, raw = get_application_status("application_1")
+    assert state == "RUNNING"
+    assert "RUNNING" in raw
+
+
+def test_status_raises_on_nonzero_return():
+    fake = _fake_proc(1, stderr="not found")
+    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
+        with pytest.raises(YarnError):
+            get_application_status("application_x")
+
+
+def test_logs_returns_stdout():
+    fake = _fake_proc(0, stdout="log line 1\nlog line 2\n")
+    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake) as m:
+        out = get_application_logs("application_1")
+    assert out == "log line 1\nlog line 2\n"
+    args = m.call_args.args[0]
+    assert args[:3] == ["yarn", "logs", "-applicationId"]
+    assert args[3] == "application_1"
+
+
+def test_kill_invokes_yarn_application_kill():
+    fake = _fake_proc(0)
+    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake) as m:
+        kill_application("application_1")
+    args = m.call_args.args[0]
+    assert args[:3] == ["yarn", "application", "-kill"]
+    assert args[3] == "application_1"
+
+
+def test_kill_raises_on_nonzero_return():
+    fake = _fake_proc(1, stderr="denied")
+    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
+        with pytest.raises(YarnError):
+            kill_application("application_1")