diff --git a/spark_executor/core/yarn_client.py b/spark_executor/core/yarn_client.py new file mode 100644 index 0000000..15bcec1 --- /dev/null +++ b/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}" + ) diff --git a/tests/unit/test_yarn_client.py b/tests/unit/test_yarn_client.py new file mode 100644 index 0000000..be92a5c --- /dev/null +++ b/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")