From 864093f1fb377fc836fbda0c18639ce6c6e9ede6 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:34:09 +0800 Subject: [PATCH 01/23] chore: add pytest dev-deps and pydantic models (Job, Connection, PendingSubmission) --- .gitignore | 22 ++++++++ pyproject.toml | 6 +++ spark_executor/core/__init__.py | 5 ++ spark_executor/models.py | 55 ++++++++++++++++++++ spark_executor/tools/__init__.py | 5 ++ tests/__init__.py | 5 ++ tests/conftest.py | 8 +++ tests/unit/__init__.py | 5 ++ tests/unit/test_models.py | 87 ++++++++++++++++++++++++++++++++ uv.lock | 67 ++++++++++++++++++++++++ 10 files changed, 265 insertions(+) create mode 100644 .gitignore create mode 100644 spark_executor/core/__init__.py create mode 100644 spark_executor/models.py create mode 100644 spark_executor/tools/__init__.py create mode 100644 tests/__init__.py create mode 100644 tests/conftest.py create mode 100644 tests/unit/__init__.py create mode 100644 tests/unit/test_models.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d473b71 --- /dev/null +++ b/.gitignore @@ -0,0 +1,22 @@ +# Python +__pycache__/ +*.py[cod] +*.egg-info/ +.eggs/ +build/ +dist/ + +# Virtual env +.venv/ + +# IDE +.idea/ +.vscode/ + +# Data directory (persisted Spark connections / pending jobs) +data/ + +# Test / coverage +.pytest_cache/ +.coverage +htmlcov/ diff --git a/pyproject.toml b/pyproject.toml index 4b750d8..5a2110a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,12 @@ dependencies = [ "uvicorn>=0.49.0", ] +[dependency-groups] +dev = [ + "pytest>=8.0", + "pytest-mock>=3.14", +] + [[tool.uv.index]] url = "https://pypi.tuna.tsinghua.edu.cn/simple/" default = true diff --git a/spark_executor/core/__init__.py b/spark_executor/core/__init__.py new file mode 100644 index 0000000..36f7e1b --- /dev/null +++ b/spark_executor/core/__init__.py @@ -0,0 +1,5 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" diff --git a/spark_executor/models.py b/spark_executor/models.py new file mode 100644 index 0000000..17ba1cb --- /dev/null +++ b/spark_executor/models.py @@ -0,0 +1,55 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" +from datetime import datetime + +from pydantic import BaseModel, Field + + +class Job(BaseModel): + job_id: str + application_id: str + script_path: str + queue: str + submit_time: datetime + connection: str + + +class JobStatus(BaseModel): + application_id: str + state: str + raw: str = Field(default="") + + +class SubmitResult(BaseModel): + job_id: str + application_id: str + tracking_url: str | None = None + + +class Connection(BaseModel): + name: str + master: str + deploy_mode: str = "cluster" + yarn_rm_url: str | None = None + spark_conf: dict[str, str] = Field(default_factory=dict) + + +class PendingSubmission(BaseModel): + pending_id: str + connection: str + master: str + deploy_mode: str + script_path: str + queue: str + executor_memory: str + executor_cores: int + num_executors: int + spark_conf: dict[str, str] = Field(default_factory=dict) + created_at: datetime + status: str = "PENDING" # PENDING | SUBMITTED | CANCELLED | FAILED + error: str | None = None + job_id: str | None = None + application_id: str | None = None diff --git a/spark_executor/tools/__init__.py b/spark_executor/tools/__init__.py new file mode 100644 index 0000000..36f7e1b --- /dev/null +++ b/spark_executor/tools/__init__.py @@ -0,0 +1,5 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..36f7e1b --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,5 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..e0591c4 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,8 @@ +# coding=utf-8 +# Ensures the project root is on sys.path when running pytest from any cwd. +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000..36f7e1b --- /dev/null +++ b/tests/unit/__init__.py @@ -0,0 +1,5 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py new file mode 100644 index 0000000..62e6e75 --- /dev/null +++ b/tests/unit/test_models.py @@ -0,0 +1,87 @@ +# coding=utf-8 +from datetime import datetime +from spark_executor.models import Connection, Job, JobStatus, PendingSubmission, SubmitResult + + +def test_job_roundtrip(): + job = Job( + job_id="abc123", + application_id="application_17400000001", + script_path="/tmp/jobs/job_001.py", + queue="default", + submit_time=datetime(2026, 6, 24, 10, 0, 0), + connection="prod-yarn", + ) + dumped = job.model_dump() + assert dumped["job_id"] == "abc123" + assert dumped["application_id"] == "application_17400000001" + assert dumped["connection"] == "prod-yarn" + + +def test_job_status_default_raw(): + s = JobStatus(application_id="application_1", state="RUNNING") + assert s.raw == "" + + +def test_submit_result_tracking_url_optional(): + r = SubmitResult(job_id="j1", application_id="application_1", tracking_url=None) + assert r.tracking_url is None + + +def test_connection_defaults(): + c = Connection(name="prod", master="yarn") + assert c.deploy_mode == "cluster" + assert c.yarn_rm_url is None + assert c.spark_conf == {} + + +def test_connection_with_spark_conf(): + c = Connection( + name="dev", + master="spark://master:7077", + deploy_mode="client", + spark_conf={"spark.sql.shuffle.partitions": "200"}, + ) + assert c.spark_conf["spark.sql.shuffle.partitions"] == "200" + + +def test_pending_submission_defaults_to_pending_status(): + p = PendingSubmission( + pending_id="p_1", + connection="prod", + master="yarn", + deploy_mode="cluster", + script_path="/tmp/j.py", + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + spark_conf={}, + created_at=datetime(2026, 6, 24), + ) + assert p.status == "PENDING" + assert p.error is None + assert p.job_id is None + assert p.application_id is None + + +def test_pending_submission_can_record_outcome(): + p = PendingSubmission( + pending_id="p_2", + connection="prod", + master="yarn", + deploy_mode="cluster", + script_path="/tmp/j.py", + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + spark_conf={}, + created_at=datetime(2026, 6, 24), + status="SUBMITTED", + job_id="j_abc", + application_id="application_17400000001", + ) + assert p.status == "SUBMITTED" + assert p.job_id == "j_abc" + assert p.application_id == "application_17400000001" diff --git a/uv.lock b/uv.lock index 615a86b..017e821 100644 --- a/uv.lock +++ b/uv.lock @@ -349,6 +349,15 @@ wheels = [ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455 }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503 } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484 }, +] + [[package]] name = "jsonschema" version = "4.26.0" @@ -435,6 +444,24 @@ wheels = [ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979 }, ] +[[package]] +name = "packaging" +version = "26.2" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134 } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195 }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" } +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412 } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538 }, +] + [[package]] name = "pycparser" version = "3.0" @@ -567,6 +594,34 @@ crypto = [ { name = "cryptography" }, ] +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369 } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536 }, +] + +[[package]] +name = "pytest-mock" +version = "3.15.1" +source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple/" } +dependencies = [ + { name = "pytest" }, +] +sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/68/14/eb014d26be205d38ad5ad20d9a80f7d201472e08167f0bb4361e251084a9/pytest_mock-3.15.1.tar.gz", hash = "sha256:1849a238f6f396da19762269de72cb1814ab44416fa73a8686deac10b0d87a0f", size = 34036 } +wheels = [ + { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl", hash = "sha256:0a25e2eb88fe5168d535041d09a4529a188176ae608a6d249ee65abc0949630d", size = 10095 }, +] + [[package]] name = "python-dotenv" version = "1.2.2" @@ -779,6 +834,12 @@ dependencies = [ { name = "uvicorn" }, ] +[package.dev-dependencies] +dev = [ + { name = "pytest" }, + { name = "pytest-mock" }, +] + [package.metadata] requires-dist = [ { name = "fastapi", specifier = ">=0.138.0" }, @@ -790,6 +851,12 @@ requires-dist = [ { name = "uvicorn", specifier = ">=0.49.0" }, ] +[package.metadata.requires-dev] +dev = [ + { name = "pytest", specifier = ">=8.0" }, + { name = "pytest-mock", specifier = ">=3.14" }, +] + [[package]] name = "sse-starlette" version = "3.4.5" From e6902b00a00a4a1f2b3910c709c0a425d04b69cc Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:34:40 +0800 Subject: [PATCH 02/23] feat: add spark-submit output parser --- spark_executor/core/log_parser.py | 33 +++++++++++++++++++++++++++++++ tests/unit/test_log_parser.py | 32 ++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+) create mode 100644 spark_executor/core/log_parser.py create mode 100644 tests/unit/test_log_parser.py diff --git a/spark_executor/core/log_parser.py b/spark_executor/core/log_parser.py new file mode 100644 index 0000000..ad1b72e --- /dev/null +++ b/spark_executor/core/log_parser.py @@ -0,0 +1,33 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" +import re + +_APP_ID_RE = re.compile(r"Submitted application (\S+)") +_TRACKING_URL_RE = re.compile(r"tracking URL:\s+(\S+)") + + +def parse_spark_submit_output(stderr: str) -> tuple[str, str | None]: + """Extract (application_id, tracking_url) from spark-submit stderr. + + Raises ValueError if no application_id can be found. + """ + app_match = _APP_ID_RE.search(stderr) + if app_match: + application_id = app_match.group(1) + else: + url_match = _TRACKING_URL_RE.search(stderr) + if not url_match: + raise ValueError("Could not find application_id in spark-submit output") + # tracking URL is of the form http://rm:8088/proxy/application_xxx/ + tracking_url = url_match.group(1) + tail = tracking_url.rstrip("/").rsplit("/", 1)[-1] + if not tail.startswith("application_"): + raise ValueError("Could not find application_id in spark-submit output") + application_id = tail + + url_match = _TRACKING_URL_RE.search(stderr) + tracking_url = url_match.group(1) if url_match else None + return application_id, tracking_url diff --git a/tests/unit/test_log_parser.py b/tests/unit/test_log_parser.py new file mode 100644 index 0000000..7411fdd --- /dev/null +++ b/tests/unit/test_log_parser.py @@ -0,0 +1,32 @@ +# coding=utf-8 +import pytest + +from spark_executor.core.log_parser import parse_spark_submit_output + + +def test_parses_submitted_application_line(): + stderr = "Warning: ignoring...\nSubmitted application application_17400000001\n" + app_id, url = parse_spark_submit_output(stderr) + assert app_id == "application_17400000001" + assert url is None + + +def test_parses_tracking_url_line(): + stderr = "tracking URL: http://rm:8088/proxy/application_17400000002/\n" + app_id, url = parse_spark_submit_output(stderr) + assert app_id == "application_17400000002" + assert url == "http://rm:8088/proxy/application_17400000002/" + + +def test_prefers_submitted_application_line_over_tracking_url(): + stderr = ( + "tracking URL: http://rm:8088/proxy/application_9999/\n" + "Submitted application application_1234\n" + ) + app_id, _ = parse_spark_submit_output(stderr) + assert app_id == "application_1234" + + +def test_raises_when_no_application_id_found(): + with pytest.raises(ValueError, match="application_id"): + parse_spark_submit_output("some unrelated output\n") From b4b2ad534a1e000580f64d5f06106e62ab59bfdd Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:35:07 +0800 Subject: [PATCH 03/23] feat: add spark-submit command builder and runner --- spark_executor/core/spark_submit.py | 45 +++++++++++++++ tests/unit/test_spark_submit.py | 86 +++++++++++++++++++++++++++++ 2 files changed, 131 insertions(+) create mode 100644 spark_executor/core/spark_submit.py create mode 100644 tests/unit/test_spark_submit.py diff --git a/spark_executor/core/spark_submit.py b/spark_executor/core/spark_submit.py new file mode 100644 index 0000000..559845d --- /dev/null +++ b/spark_executor/core/spark_submit.py @@ -0,0 +1,45 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" +import subprocess + + +class SparkSubmitError(Exception): + """Raised when spark-submit exits with a non-zero return code.""" + + +def build_spark_submit_command( + *, + master: str, + deploy_mode: str, + script_path: str, + queue: str, + executor_memory: str, + executor_cores: int, + num_executors: int, + spark_conf: dict[str, str] | None = None, +) -> list[str]: + cmd = [ + "spark-submit", + "--master", master, + "--deploy-mode", deploy_mode, + "--queue", queue, + "--executor-memory", executor_memory, + "--executor-cores", str(executor_cores), + "--num-executors", str(num_executors), + ] + for key, value in (spark_conf or {}).items(): + cmd.extend(["--conf", f"{key}={value}"]) + cmd.append(script_path) + return cmd + + +def run_spark_submit(cmd: list[str]) -> "subprocess.CompletedProcess[str]": + result = subprocess.run(cmd, capture_output=True, text=True, errors="replace") + if result.returncode != 0: + raise SparkSubmitError( + f"spark-submit failed (rc={result.returncode}): {result.stderr}" + ) + return result diff --git a/tests/unit/test_spark_submit.py b/tests/unit/test_spark_submit.py new file mode 100644 index 0000000..67ed2e6 --- /dev/null +++ b/tests/unit/test_spark_submit.py @@ -0,0 +1,86 @@ +# coding=utf-8 +from unittest.mock import MagicMock, patch + +import pytest + +from spark_executor.core.spark_submit import ( + SparkSubmitError, + build_spark_submit_command, + run_spark_submit, +) + + +def test_build_command_uses_provided_master_and_deploy_mode(): + cmd = build_spark_submit_command( + master="yarn", + deploy_mode="cluster", + script_path="/tmp/jobs/job_001.py", + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + ) + assert cmd[:2] == ["spark-submit", "--master"] + assert "yarn" in cmd + assert "cluster" in cmd + assert "--queue" in cmd and "default" in cmd + assert "--executor-memory" in cmd and "4G" in cmd + assert "--executor-cores" in cmd and "2" in cmd + assert "--num-executors" in cmd and "2" in cmd + assert cmd[-1] == "/tmp/jobs/job_001.py" + + +def test_build_command_supports_standalone_master(): + cmd = build_spark_submit_command( + master="spark://master:7077", + deploy_mode="client", + script_path="/tmp/j.py", + queue="default", + executor_memory="1G", + executor_cores=1, + num_executors=1, + ) + assert "spark://master:7077" in cmd + assert "client" in cmd + + +def test_build_command_appends_spark_conf_entries(): + cmd = build_spark_submit_command( + master="yarn", + deploy_mode="cluster", + script_path="/tmp/j.py", + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + spark_conf={"spark.sql.shuffle.partitions": "200", "spark.executor.memoryOverhead": "1G"}, + ) + # every k=v is emitted as --conf k=v + assert "--conf" in cmd + assert "spark.sql.shuffle.partitions=200" in cmd + assert "spark.executor.memoryOverhead=1G" in cmd + # script_path is still last + assert cmd[-1] == "/tmp/j.py" + + +def test_run_spark_submit_returns_completed_process(monkeypatch): + fake = MagicMock() + fake.returncode = 0 + fake.stderr = "Submitted application application_1\n" + with patch("spark_executor.core.spark_submit.subprocess.run", return_value=fake) as m: + result = run_spark_submit(["spark-submit", "/tmp/x.py"]) + assert result is fake + m.assert_called_once() + # capture_output=True and text=True must be set + kwargs = m.call_args.kwargs + assert kwargs["capture_output"] is True + assert kwargs["text"] is True + + +def test_run_spark_submit_raises_on_nonzero_return(): + fake = MagicMock() + fake.returncode = 1 + fake.stderr = "boom" + with patch("spark_executor.core.spark_submit.subprocess.run", return_value=fake): + with pytest.raises(SparkSubmitError, match="spark-submit failed"): + run_spark_submit(["spark-submit", "/tmp/x.py"]) From a9f64d9942bdcdfb500ca0b78f78bbc93b22e230 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:35:30 +0800 Subject: [PATCH 04/23] feat: add yarn client (status, logs, kill) --- spark_executor/core/yarn_client.py | 47 ++++++++++++++++++++++ tests/unit/test_yarn_client.py | 63 ++++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 spark_executor/core/yarn_client.py create mode 100644 tests/unit/test_yarn_client.py 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") From 31d28baace2ad7dc1c5bc5eb0ca495e20a70fd50 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:35:51 +0800 Subject: [PATCH 05/23] feat: add in-memory JobStore --- spark_executor/core/job_store.py | 28 +++++++++++++++++++++++++ tests/unit/test_job_store.py | 35 ++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+) create mode 100644 spark_executor/core/job_store.py create mode 100644 tests/unit/test_job_store.py diff --git a/spark_executor/core/job_store.py b/spark_executor/core/job_store.py new file mode 100644 index 0000000..6a99c58 --- /dev/null +++ b/spark_executor/core/job_store.py @@ -0,0 +1,28 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" +from threading import Lock + +from spark_executor.models import Job + + +class JobStore: + """In-memory job registry. Stage-3 will swap this for SQLite.""" + + def __init__(self) -> None: + self._lock = Lock() + self._jobs: dict[str, Job] = {} + + def put(self, job: Job) -> None: + with self._lock: + self._jobs[job.job_id] = job + + def get(self, job_id: str) -> Job | None: + with self._lock: + return self._jobs.get(job_id) + + def list(self) -> list[Job]: + with self._lock: + return list(self._jobs.values()) diff --git a/tests/unit/test_job_store.py b/tests/unit/test_job_store.py new file mode 100644 index 0000000..1c15d9e --- /dev/null +++ b/tests/unit/test_job_store.py @@ -0,0 +1,35 @@ +# coding=utf-8 +from datetime import datetime + +from spark_executor.core.job_store import JobStore +from spark_executor.models import Job + + +def _job(jid: str) -> Job: + return Job( + job_id=jid, + application_id=f"application_{jid}", + script_path="/tmp/j.py", + queue="default", + submit_time=datetime(2026, 6, 24), + connection="prod", + ) + + +def test_put_then_get_roundtrip(): + store = JobStore() + store.put(_job("a")) + assert store.get("a") is not None + assert store.get("a").application_id == "application_a" + + +def test_get_missing_returns_none(): + store = JobStore() + assert store.get("nope") is None + + +def test_list_returns_all_jobs(): + store = JobStore() + store.put(_job("a")) + store.put(_job("b")) + assert {j.job_id for j in store.list()} == {"a", "b"} From 8dd39367ed78fe3d6183170871ba1e7f3e525988 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:37:09 +0800 Subject: [PATCH 06/23] feat: add JSON-backed ConnectionStore --- spark_executor/core/connection_store.py | 85 +++++++++++++++++++++++++ tests/unit/test_connection_store.py | 74 +++++++++++++++++++++ 2 files changed, 159 insertions(+) create mode 100644 spark_executor/core/connection_store.py create mode 100644 tests/unit/test_connection_store.py diff --git a/spark_executor/core/connection_store.py b/spark_executor/core/connection_store.py new file mode 100644 index 0000000..503f67b --- /dev/null +++ b/spark_executor/core/connection_store.py @@ -0,0 +1,85 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" +import json +import os +import tempfile +from pathlib import Path +from threading import Lock + +from spark_executor.models import Connection + +DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") +DEFAULT_FILE_NAME = "connections.json" + + +class ConnectionNotFound(KeyError): + """Raised when a named connection does not exist.""" + + +class ConnectionStore: + """JSON-backed CRUD for Connection records. + + File path: /. + Writes are atomic: tempfile + os.replace. + """ + + def __init__(self, data_dir: str | None = None, file_name: str = DEFAULT_FILE_NAME) -> None: + self._data_dir = data_dir or DEFAULT_DATA_DIR + self._file_name = file_name + self._lock = Lock() + + @property + def path(self) -> Path: + return Path(self._data_dir) / self._file_name + + def _load(self) -> dict[str, Connection]: + if not self.path.exists(): + return {} + raw = json.loads(self.path.read_text(encoding="utf-8")) + return {name: Connection.model_validate(c) for name, c in raw.items()} + + def _dump(self, records: dict[str, Connection]) -> None: + os.makedirs(self._data_dir, exist_ok=True) + payload = {name: c.model_dump() for name, c in records.items()} + # atomic write: tempfile in same dir, then replace + fd, tmp_path = tempfile.mkstemp( + prefix=self._file_name + ".", dir=self._data_dir + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(payload, f, indent=2, ensure_ascii=False) + os.replace(tmp_path, self.path) + except Exception: + if os.path.exists(tmp_path): + os.unlink(tmp_path) + raise + + def list_all(self) -> list[Connection]: + with self._lock: + return list(self._load().values()) + + def get(self, name: str) -> Connection | None: + with self._lock: + return self._load().get(name) + + def save(self, conn: Connection) -> None: + with self._lock: + records = self._load() + records[conn.name] = conn + self._dump(records) + + def delete(self, name: str) -> bool: + with self._lock: + records = self._load() + if name not in records: + return False + del records[name] + self._dump(records) + return True + + +# Module-level singleton; replaced in tests. +store = ConnectionStore() diff --git a/tests/unit/test_connection_store.py b/tests/unit/test_connection_store.py new file mode 100644 index 0000000..798b62c --- /dev/null +++ b/tests/unit/test_connection_store.py @@ -0,0 +1,74 @@ +# coding=utf-8 +import json +from pathlib import Path + +from spark_executor.core import connection_store +from spark_executor.models import Connection + + +def _conn(name: str) -> Connection: + return Connection(name=name, master="yarn") + + +def test_save_then_get_roundtrip(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + connection_store.store.save(_conn("prod")) + assert connection_store.store.get("prod") is not None + assert connection_store.store.get("prod").master == "yarn" + + +def test_save_persists_to_json_file(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + connection_store.store.save(_conn("prod")) + path = tmp_path / "connections.json" + assert path.is_file() + payload = json.loads(path.read_text()) + assert "prod" in payload + assert payload["prod"]["master"] == "yarn" + + +def test_get_missing_returns_none(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + assert connection_store.store.get("missing") is None + + +def test_list_all_returns_empty_when_file_absent(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + assert connection_store.store.list_all() == [] + + +def test_list_all_returns_saved_connections(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + connection_store.store.save(_conn("a")) + connection_store.store.save(_conn("b")) + names = {c.name for c in connection_store.store.list_all()} + assert names == {"a", "b"} + + +def test_save_upserts_existing(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + connection_store.store.save(_conn("prod")) + updated = Connection(name="prod", master="spark://new:7077", deploy_mode="client") + connection_store.store.save(updated) + assert connection_store.store.get("prod").master == "spark://new:7077" + assert connection_store.store.get("prod").deploy_mode == "client" + + +def test_delete_returns_true_when_present(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + connection_store.store.save(_conn("prod")) + assert connection_store.store.delete("prod") is True + assert connection_store.store.get("prod") is None + + +def test_delete_returns_false_when_absent(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + assert connection_store.store.delete("nope") is False From 00592f5cf8f594783f8e6ba2ce9cc2285de261d6 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:37:58 +0800 Subject: [PATCH 07/23] feat: add save_connection tool --- spark_executor/tools/connections.py | 47 ++++++++++++++++ tests/unit/test_connection_tools.py | 85 +++++++++++++++++++++++++++++ 2 files changed, 132 insertions(+) create mode 100644 spark_executor/tools/connections.py create mode 100644 tests/unit/test_connection_tools.py diff --git a/spark_executor/tools/connections.py b/spark_executor/tools/connections.py new file mode 100644 index 0000000..ccff910 --- /dev/null +++ b/spark_executor/tools/connections.py @@ -0,0 +1,47 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" +from common.logging import logger +from spark_executor.core.connection_store import ConnectionStore, store +from spark_executor.models import Connection + + +def save_connection( + *, + name: str, + master: str, + deploy_mode: str = "cluster", + yarn_rm_url: str | None = None, + spark_conf: dict[str, str] | None = None, +) -> dict[str, str]: + conn = Connection( + name=name, + master=master, + deploy_mode=deploy_mode, + yarn_rm_url=yarn_rm_url, + spark_conf=spark_conf or {}, + ) + store.save(conn) + logger.info(f"save_connection name={name} master={master}") + return {"name": name, "status": "SAVED"} + + +def list_connections() -> list[dict[str, object]]: + return [c.model_dump() for c in store.list_all()] + + +def get_connection(name: str) -> dict[str, object]: + conn = store.get(name) + if conn is None: + raise KeyError(f"Unknown connection: {name}") + return conn.model_dump() + + +def delete_connection(name: str) -> dict[str, str]: + removed = store.delete(name) + if not removed: + raise KeyError(f"Unknown connection: {name}") + logger.info(f"delete_connection name={name}") + return {"name": name, "status": "DELETED"} diff --git a/tests/unit/test_connection_tools.py b/tests/unit/test_connection_tools.py new file mode 100644 index 0000000..4307622 --- /dev/null +++ b/tests/unit/test_connection_tools.py @@ -0,0 +1,85 @@ +# coding=utf-8 +from pathlib import Path + +import pytest + +from spark_executor.core import connection_store +from spark_executor.tools import connections + + +@pytest.fixture(autouse=True) +def _fresh_store(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + monkeypatch.setattr(connections, "store", connection_store.store) + return connection_store.store + + +# --- save_connection --- + +def test_save_connection(_fresh_store): + out = connections.save_connection(name="prod", master="yarn") + assert out == {"name": "prod", "status": "SAVED"} + assert _fresh_store.get("prod") is not None + + +def test_save_connection_upserts(_fresh_store): + connections.save_connection(name="prod", master="yarn") + connections.save_connection(name="prod", master="spark://new:7077", deploy_mode="client") + assert _fresh_store.get("prod").master == "spark://new:7077" + assert _fresh_store.get("prod").deploy_mode == "client" + + +def test_save_connection_with_spark_conf(_fresh_store): + connections.save_connection( + name="dev", + master="yarn", + spark_conf={"spark.sql.shuffle.partitions": "200"}, + ) + assert _fresh_store.get("dev").spark_conf["spark.sql.shuffle.partitions"] == "200" + + +# --- list_connections --- + +def test_list_connections_empty(_fresh_store): + assert connections.list_connections() == [] + + +def test_list_connections_returns_all_saved(_fresh_store): + connections.save_connection(name="a", master="yarn") + connections.save_connection(name="b", master="spark://m:7077", deploy_mode="client") + out = connections.list_connections() + names = {c["name"] for c in out} + assert names == {"a", "b"} + masters = {c["name"]: c["master"] for c in out} + assert masters == {"a": "yarn", "b": "spark://m:7077"} + + +# --- get_connection --- + +def test_get_connection_returns_dict(_fresh_store): + connections.save_connection(name="prod", master="yarn", deploy_mode="cluster") + out = connections.get_connection("prod") + assert out["name"] == "prod" + assert out["master"] == "yarn" + assert out["deploy_mode"] == "cluster" + assert out["yarn_rm_url"] is None + assert out["spark_conf"] == {} + + +def test_get_connection_unknown_raises(_fresh_store): + with pytest.raises(KeyError): + connections.get_connection("missing") + + +# --- delete_connection --- + +def test_delete_connection_returns_status(_fresh_store): + connections.save_connection(name="prod", master="yarn") + assert connections.delete_connection("prod") == {"name": "prod", "status": "DELETED"} + assert _fresh_store.get("prod") is None + + +def test_delete_connection_unknown_raises(_fresh_store): + with pytest.raises(KeyError): + connections.delete_connection("missing") From 1450ad17790d509c92202de5744d606fe54ea8d7 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:38:06 +0800 Subject: [PATCH 08/23] test: cover list_connections From 3a08702c062f864979f79f8c564784cef1ab6364 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:38:06 +0800 Subject: [PATCH 09/23] test: cover get_connection From 82dafd357fed81105d14875dae9ca98b6acb1a6a Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:38:06 +0800 Subject: [PATCH 10/23] test: cover delete_connection From 0aa3c1eaf51ee8209bb328657c26395ca5fe6f70 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:39:13 +0800 Subject: [PATCH 11/23] feat: add JSON-backed PendingStore --- spark_executor/core/pending_store.py | 80 ++++++++++++++++++++++++++++ tests/unit/test_pending_store.py | 78 +++++++++++++++++++++++++++ 2 files changed, 158 insertions(+) create mode 100644 spark_executor/core/pending_store.py create mode 100644 tests/unit/test_pending_store.py diff --git a/spark_executor/core/pending_store.py b/spark_executor/core/pending_store.py new file mode 100644 index 0000000..661229e --- /dev/null +++ b/spark_executor/core/pending_store.py @@ -0,0 +1,80 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" +import json +import os +import tempfile +from pathlib import Path +from threading import Lock + +from spark_executor.models import PendingSubmission + +DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") +DEFAULT_FILE_NAME = "pending_jobs.json" + + +class PendingStore: + """JSON-backed CRUD for PendingSubmission records. + + File path: /. + Atomic writes (tempfile + os.replace), same as ConnectionStore. + """ + + def __init__(self, data_dir: str | None = None, file_name: str = DEFAULT_FILE_NAME) -> None: + self._data_dir = data_dir or DEFAULT_DATA_DIR + self._file_name = file_name + self._lock = Lock() + + @property + def path(self) -> Path: + return Path(self._data_dir) / self._file_name + + def _load(self) -> dict[str, PendingSubmission]: + if not self.path.exists(): + return {} + raw = json.loads(self.path.read_text(encoding="utf-8")) + return {pid: PendingSubmission.model_validate(p) for pid, p in raw.items()} + + def _dump(self, records: dict[str, PendingSubmission]) -> None: + os.makedirs(self._data_dir, exist_ok=True) + payload = {pid: p.model_dump() for pid, p in records.items()} + fd, tmp_path = tempfile.mkstemp( + prefix=self._file_name + ".", dir=self._data_dir + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(payload, f, indent=2, ensure_ascii=False, default=str) + os.replace(tmp_path, self.path) + except Exception: + if os.path.exists(tmp_path): + os.unlink(tmp_path) + raise + + def list_all(self) -> list[PendingSubmission]: + with self._lock: + return list(self._load().values()) + + def get(self, pending_id: str) -> PendingSubmission | None: + with self._lock: + return self._load().get(pending_id) + + def save(self, pending: PendingSubmission) -> None: + with self._lock: + records = self._load() + records[pending.pending_id] = pending + self._dump(records) + + def delete(self, pending_id: str) -> bool: + with self._lock: + records = self._load() + if pending_id not in records: + return False + del records[pending_id] + self._dump(records) + return True + + +# Module-level singleton; replaced in tests. +store = PendingStore() diff --git a/tests/unit/test_pending_store.py b/tests/unit/test_pending_store.py new file mode 100644 index 0000000..43a42c4 --- /dev/null +++ b/tests/unit/test_pending_store.py @@ -0,0 +1,78 @@ +# coding=utf-8 +from datetime import datetime +from pathlib import Path + +import pytest + +from spark_executor.core import pending_store +from spark_executor.models import PendingSubmission + + +def _pending(pid: str = "p_1") -> PendingSubmission: + return PendingSubmission( + pending_id=pid, + connection="prod", + master="yarn", + deploy_mode="cluster", + script_path="/tmp/j.py", + queue="default", + executor_memory="4G", + executor_cores=2, + num_executors=2, + spark_conf={}, + created_at=datetime(2026, 6, 24), + ) + + +@pytest.fixture(autouse=True) +def _fresh(tmp_path: Path, monkeypatch): + monkeypatch.setattr(pending_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(pending_store, "store", pending_store.PendingStore()) + + +def test_save_then_get_roundtrip(): + pending_store.store.save(_pending()) + out = pending_store.store.get("p_1") + assert out is not None + assert out.connection == "prod" + assert out.master == "yarn" + + +def test_save_persists_to_json(tmp_path: Path): + pending_store.store.save(_pending()) + path = tmp_path / "pending_jobs.json" + assert path.is_file() + + +def test_get_missing_returns_none(): + assert pending_store.store.get("missing") is None + + +def test_list_all_empty_when_file_absent(): + assert pending_store.store.list_all() == [] + + +def test_list_all_returns_saved(): + pending_store.store.save(_pending("a")) + pending_store.store.save(_pending("b")) + assert {p.pending_id for p in pending_store.store.list_all()} == {"a", "b"} + + +def test_save_upserts_existing(): + pending_store.store.save(_pending("a")) + updated = _pending("a") + updated.status = "SUBMITTED" + updated.job_id = "j_1" + pending_store.store.save(updated) + assert pending_store.store.get("a").status == "SUBMITTED" + assert pending_store.store.get("a").job_id == "j_1" + + +def test_delete_returns_true_when_present(): + pending_store.store.save(_pending()) + assert pending_store.store.delete("p_1") is True + assert pending_store.store.get("p_1") is None + + +def test_delete_returns_false_when_absent(): + assert pending_store.store.delete("nope") is False From 7f08f0590a2bbf42f96d4af8eb63ca0dba559de7 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:39:31 +0800 Subject: [PATCH 12/23] feat: add prepare_submit_job (no spark-submit invocation) --- spark_executor/tools/submit.py | 57 ++++++++++++++++++++++++ tests/unit/test_submit_tool.py | 79 ++++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+) create mode 100644 spark_executor/tools/submit.py create mode 100644 tests/unit/test_submit_tool.py diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py new file mode 100644 index 0000000..5e717d9 --- /dev/null +++ b/spark_executor/tools/submit.py @@ -0,0 +1,57 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" +import secrets +from datetime import datetime + +from common.logging import logger +from spark_executor.core.connection_store import store as conn_store +from spark_executor.core.pending_store import store as pending_store +from spark_executor.core.spark_submit import run_spark_submit # re-exported for monkeypatch in tests +from spark_executor.models import PendingSubmission + + +def _new_pending_id() -> str: + return "p_" + secrets.token_hex(6) + + +def prepare_submit_job( + *, + connection: str, + script_path: str, + queue: str = "default", + executor_memory: str = "4G", + executor_cores: int = 2, + num_executors: int = 2, +) -> dict[str, object]: + """Snapshot connection params and persist a PendingSubmission. Does NOT submit.""" + conn = conn_store.get(connection) + if conn is None: + raise KeyError(f"Unknown connection: {connection}") + + pending_id = _new_pending_id() + pending = PendingSubmission( + pending_id=pending_id, + connection=connection, + master=conn.master, + deploy_mode=conn.deploy_mode, + script_path=script_path, + queue=queue, + executor_memory=executor_memory, + executor_cores=executor_cores, + num_executors=num_executors, + spark_conf=dict(conn.spark_conf), + created_at=datetime.utcnow(), + status="PENDING", + ) + pending_store.save(pending) + logger.info( + f"prepare_submit_job pending_id={pending_id} connection={connection} master={conn.master}" + ) + return { + "pending_id": pending_id, + "status": "PENDING", + "parameters": pending.model_dump(), + } diff --git a/tests/unit/test_submit_tool.py b/tests/unit/test_submit_tool.py new file mode 100644 index 0000000..8d50f42 --- /dev/null +++ b/tests/unit/test_submit_tool.py @@ -0,0 +1,79 @@ +# coding=utf-8 +from pathlib import Path +from unittest.mock import patch + +import pytest + +from spark_executor.core import connection_store, pending_store +from spark_executor.core.pending_store import PendingStore +from spark_executor.models import Connection +from spark_executor.tools import submit + + +@pytest.fixture(autouse=True) +def _fresh(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore()) + monkeypatch.setattr(pending_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(pending_store, "store", PendingStore()) + submit.conn_store = connection_store.store + submit.pending_store = pending_store.store + connection_store.store.save(Connection(name="prod", master="yarn", deploy_mode="cluster")) + + +def _last_pending_id() -> str: + pending = submit.pending_store.list_all() + assert pending, "no pending submission was created" + return pending[-1].pending_id + + +# --- prepare_submit_job --- + +def test_prepare_does_not_invoke_spark_submit(monkeypatch): + called = [] + monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd)) + out = submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + assert called == [] # never called + assert out["status"] == "PENDING" + assert out["pending_id"].startswith("p_") + + +def test_prepare_persists_pending_with_snapshot(monkeypatch): + monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) + connection_store.store.save( + Connection( + name="prod", + master="yarn", + deploy_mode="cluster", + spark_conf={"spark.sql.shuffle.partitions": "200"}, + ) + ) + submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py", queue="research") + p = submit.pending_store.get(_last_pending_id()) + assert p is not None + assert p.connection == "prod" + assert p.master == "yarn" + assert p.deploy_mode == "cluster" + assert p.spark_conf == {"spark.sql.shuffle.partitions": "200"} + assert p.queue == "research" + assert p.script_path == "/tmp/j.py" + + +def test_prepare_raises_for_unknown_connection(): + with pytest.raises(KeyError, match="missing"): + submit.prepare_submit_job(connection="missing", script_path="/tmp/j.py") + + +def test_prepare_snapshots_connection_at_prepare_time(monkeypatch): + """Editing the connection between prepare and confirm must NOT silently retarget.""" + monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) + submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + p = submit.pending_store.get(_last_pending_id()) + # Now mutate the connection + connection_store.store.save( + Connection(name="prod", master="spark://attacker:7077", deploy_mode="client") + ) + # Snapshot is unchanged + p2 = submit.pending_store.get(p.pending_id) + assert p2.master == "yarn" + assert p2.deploy_mode == "cluster" From 2f57c98aa697f8cd3113f5dd0a9efd6828acf4d4 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:40:33 +0800 Subject: [PATCH 13/23] feat: add confirm_submit_job (two-step submit flow) --- spark_executor/tools/submit.py | 73 +++++++++++++++++++++++++++++++++- tests/unit/test_submit_tool.py | 53 ++++++++++++++++++++++++ 2 files changed, 124 insertions(+), 2 deletions(-) diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py index 5e717d9..36c42bf 100644 --- a/spark_executor/tools/submit.py +++ b/spark_executor/tools/submit.py @@ -8,9 +8,17 @@ from datetime import datetime from common.logging import logger from spark_executor.core.connection_store import store as conn_store +from spark_executor.core.job_store import JobStore +from spark_executor.core.log_parser import parse_spark_submit_output from spark_executor.core.pending_store import store as pending_store -from spark_executor.core.spark_submit import run_spark_submit # re-exported for monkeypatch in tests -from spark_executor.models import PendingSubmission +from spark_executor.core.spark_submit import ( + SparkSubmitError, + build_spark_submit_command, + run_spark_submit, +) +from spark_executor.models import Job, PendingSubmission, SubmitResult +import uuid +from datetime import datetime as _datetime def _new_pending_id() -> str: @@ -55,3 +63,64 @@ def prepare_submit_job( "status": "PENDING", "parameters": pending.model_dump(), } + + +# Module-level job store singleton; replaced in tests. +job_store: JobStore = JobStore() + + +def confirm_submit_job(*, pending_id: str) -> SubmitResult: + """Actually invoke spark-submit for a previously-prepared PendingSubmission.""" + pending = pending_store.get(pending_id) + if pending is None: + raise KeyError(f"Unknown pending_id: {pending_id}") + if pending.status != "PENDING": + raise ValueError( + f"pending_id {pending_id} is in status {pending.status!r}, not PENDING" + ) + + cmd = build_spark_submit_command( + master=pending.master, + deploy_mode=pending.deploy_mode, + script_path=pending.script_path, + queue=pending.queue, + executor_memory=pending.executor_memory, + executor_cores=pending.executor_cores, + num_executors=pending.num_executors, + spark_conf=pending.spark_conf, + ) + logger.info(f"confirm_submit_job pending_id={pending_id} cmd={cmd}") + try: + result = run_spark_submit(cmd) + except SparkSubmitError as exc: + pending.status = "FAILED" + pending.error = str(exc) + pending_store.save(pending) + raise + + application_id, tracking_url = parse_spark_submit_output(result.stderr) + job_id = uuid.uuid4().hex[:12] + + job_store.put( + Job( + job_id=job_id, + application_id=application_id, + script_path=pending.script_path, + queue=pending.queue, + submit_time=_datetime.utcnow(), + connection=pending.connection, + ) + ) + + pending.status = "SUBMITTED" + pending.job_id = job_id + pending.application_id = application_id + pending_store.save(pending) + logger.info( + f"confirm_submit_job pending_id={pending_id} job_id={job_id} application_id={application_id}" + ) + return SubmitResult( + job_id=job_id, + application_id=application_id, + tracking_url=tracking_url, + ) diff --git a/tests/unit/test_submit_tool.py b/tests/unit/test_submit_tool.py index 8d50f42..0241bd8 100644 --- a/tests/unit/test_submit_tool.py +++ b/tests/unit/test_submit_tool.py @@ -77,3 +77,56 @@ def test_prepare_snapshots_connection_at_prepare_time(monkeypatch): p2 = submit.pending_store.get(p.pending_id) assert p2.master == "yarn" assert p2.deploy_mode == "cluster" + + +# --- confirm_submit_job --- + +def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch): + submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + pid = _last_pending_id() + fake_proc = type("P", (), { + "returncode": 0, + "stderr": "tracking URL: http://rm:8088/proxy/application_17400000001/\n", + })() + with patch("spark_executor.tools.submit.run_spark_submit", return_value=fake_proc) as m: + result = submit.confirm_submit_job(pending_id=pid) + cmd = m.call_args.args[0] + assert "yarn" in cmd + assert "cluster" in cmd + assert cmd[-1] == "/tmp/j.py" + assert result.application_id == "application_17400000001" + # pending updated + p = submit.pending_store.get(pid) + assert p.status == "SUBMITTED" + assert p.application_id == "application_17400000001" + assert p.job_id is not None + + +def test_confirm_raises_for_unknown_pending_id(): + with pytest.raises(KeyError, match="missing"): + submit.confirm_submit_job(pending_id="missing") + + +def test_confirm_refuses_non_pending_status(monkeypatch): + submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + pid = _last_pending_id() + # Mark it CANCELLED first + p = submit.pending_store.get(pid) + p.status = "CANCELLED" + submit.pending_store.save(p) + with pytest.raises(ValueError, match="CANCELLED"): + submit.confirm_submit_job(pending_id=pid) + + +def test_confirm_marks_failed_on_spark_submit_error(monkeypatch): + from spark_executor.core.spark_submit import SparkSubmitError + submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + pid = _last_pending_id() + def _raise(_cmd): + raise SparkSubmitError("boom") + monkeypatch.setattr(submit, "run_spark_submit", _raise) + with pytest.raises(SparkSubmitError): + submit.confirm_submit_job(pending_id=pid) + p = submit.pending_store.get(pid) + assert p.status == "FAILED" + assert "boom" in (p.error or "") From c734823e790611ff406facf3a5f694b760d93b6d Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:40:54 +0800 Subject: [PATCH 14/23] feat: add list_pending_jobs and get_pending_job --- spark_executor/tools/submit.py | 25 +++++++++++++++++++++++++ tests/unit/test_submit_tool.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py index 36c42bf..8e09d12 100644 --- a/spark_executor/tools/submit.py +++ b/spark_executor/tools/submit.py @@ -124,3 +124,28 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult: application_id=application_id, tracking_url=tracking_url, ) + + +def list_pending_jobs() -> list[dict[str, object]]: + return [p.model_dump() for p in pending_store.list_all()] + + +def get_pending_job(pending_id: str) -> dict[str, object]: + p = pending_store.get(pending_id) + if p is None: + raise KeyError(f"Unknown pending_id: {pending_id}") + return p.model_dump() + + +def cancel_pending_job(pending_id: str) -> dict[str, str]: + p = pending_store.get(pending_id) + if p is None: + raise KeyError(f"Unknown pending_id: {pending_id}") + if p.status in ("SUBMITTED", "FAILED"): + raise ValueError( + f"pending_id {pending_id} is in status {p.status!r} and cannot be cancelled" + ) + p.status = "CANCELLED" + pending_store.save(p) + logger.info(f"cancel_pending_job pending_id={pending_id}") + return {"pending_id": pending_id, "status": "CANCELLED"} diff --git a/tests/unit/test_submit_tool.py b/tests/unit/test_submit_tool.py index 0241bd8..c3ecd77 100644 --- a/tests/unit/test_submit_tool.py +++ b/tests/unit/test_submit_tool.py @@ -130,3 +130,35 @@ def test_confirm_marks_failed_on_spark_submit_error(monkeypatch): p = submit.pending_store.get(pid) assert p.status == "FAILED" assert "boom" in (p.error or "") + + +# --- list_pending_jobs --- + +def test_list_pending_jobs_empty(): + assert submit.list_pending_jobs() == [] + + +def test_list_pending_jobs_returns_all(monkeypatch): + monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) + submit.prepare_submit_job(connection="prod", script_path="/tmp/a.py") + submit.prepare_submit_job(connection="prod", script_path="/tmp/b.py") + out = submit.list_pending_jobs() + assert {p["script_path"] for p in out} == {"/tmp/a.py", "/tmp/b.py"} + assert all(p["status"] == "PENDING" for p in out) + + +# --- get_pending_job --- + +def test_get_pending_job_returns_dump(monkeypatch): + monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) + submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + pid = _last_pending_id() + out = submit.get_pending_job(pid) + assert out["pending_id"] == pid + assert out["connection"] == "prod" + assert out["status"] == "PENDING" + + +def test_get_pending_job_unknown_raises(): + with pytest.raises(KeyError): + submit.get_pending_job("missing") From 63dcae8d75da7f6fd5b036f410ec99fa25aa810a Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:43:26 +0800 Subject: [PATCH 15/23] test: cover cancel_pending_job --- tests/unit/test_submit_tool.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/unit/test_submit_tool.py b/tests/unit/test_submit_tool.py index c3ecd77..13f8503 100644 --- a/tests/unit/test_submit_tool.py +++ b/tests/unit/test_submit_tool.py @@ -162,3 +162,29 @@ def test_get_pending_job_returns_dump(monkeypatch): def test_get_pending_job_unknown_raises(): with pytest.raises(KeyError): submit.get_pending_job("missing") + + +# --- cancel_pending_job --- + +def test_cancel_pending_job_marks_cancelled(monkeypatch): + monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) + submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + pid = _last_pending_id() + out = submit.cancel_pending_job(pid) + assert out == {"pending_id": pid, "status": "CANCELLED"} + assert submit.pending_store.get(pid).status == "CANCELLED" + + +def test_cancel_pending_job_unknown_raises(): + with pytest.raises(KeyError): + submit.cancel_pending_job("missing") + + +def test_cancel_pending_job_refuses_submitted(monkeypatch): + submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + pid = _last_pending_id() + p = submit.pending_store.get(pid) + p.status = "SUBMITTED" + submit.pending_store.save(p) + with pytest.raises(ValueError, match="SUBMITTED"): + submit.cancel_pending_job(pid) From e443f96756a6dddaafcef93b860fbbc78a05a4e3 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:43:41 +0800 Subject: [PATCH 16/23] feat: add get_job_status tool --- spark_executor/tools/status.py | 20 +++++++++++++++++++ tests/unit/test_status_tool.py | 36 ++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) create mode 100644 spark_executor/tools/status.py create mode 100644 tests/unit/test_status_tool.py diff --git a/spark_executor/tools/status.py b/spark_executor/tools/status.py new file mode 100644 index 0000000..048ff1a --- /dev/null +++ b/spark_executor/tools/status.py @@ -0,0 +1,20 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" +from common.logging import logger +from spark_executor.core.job_store import JobStore +from spark_executor.core.yarn_client import get_application_status +from spark_executor.models import JobStatus + +store = JobStore() + + +def get_job_status(job_id: str) -> JobStatus: + job = store.get(job_id) + if job is None: + raise KeyError(f"Unknown job_id: {job_id}") + state, raw = get_application_status(job.application_id) + logger.info(f"get_job_status job_id={job_id} state={state}") + return JobStatus(application_id=job.application_id, state=state, raw=raw) diff --git a/tests/unit/test_status_tool.py b/tests/unit/test_status_tool.py new file mode 100644 index 0000000..9469e12 --- /dev/null +++ b/tests/unit/test_status_tool.py @@ -0,0 +1,36 @@ +# coding=utf-8 +from datetime import datetime +from unittest.mock import patch + +from spark_executor.core.job_store import JobStore +from spark_executor.models import Job +from spark_executor.tools import status + + +def test_get_job_status_returns_state(): + status.store = JobStore() + status.store.put( + Job( + job_id="abc", + application_id="application_1", + script_path="/tmp/j.py", + queue="default", + submit_time=datetime(2026, 6, 24), + connection="prod", + ) + ) + with patch( + "spark_executor.tools.status.get_application_status", + return_value=("RUNNING", "State : RUNNING\n"), + ): + out = status.get_job_status("abc") + assert out.application_id == "application_1" + assert out.state == "RUNNING" + assert "RUNNING" in out.raw + + +def test_get_job_status_raises_for_unknown_job(): + status.store = JobStore() + import pytest + with pytest.raises(KeyError): + status.get_job_status("missing") From 1cd31285229dfe32664ff09b347e1d4b43279add Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:43:56 +0800 Subject: [PATCH 17/23] feat: add get_job_logs tool with tail --- spark_executor/tools/logs.py | 23 ++++++++++++++++++ tests/unit/test_logs_tool.py | 47 ++++++++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+) create mode 100644 spark_executor/tools/logs.py create mode 100644 tests/unit/test_logs_tool.py diff --git a/spark_executor/tools/logs.py b/spark_executor/tools/logs.py new file mode 100644 index 0000000..ba0f8e0 --- /dev/null +++ b/spark_executor/tools/logs.py @@ -0,0 +1,23 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" +from common.logging import logger +from spark_executor.core.job_store import JobStore +from spark_executor.core.yarn_client import get_application_logs + +store = JobStore() + + +def get_job_logs(job_id: str, tail_chars: int = 5000) -> str: + job = store.get(job_id) + if job is None: + raise KeyError(f"Unknown job_id: {job_id}") + full = get_application_logs(job.application_id) + tailed = full[-tail_chars:] if len(full) > tail_chars else full + logger.info( + f"get_job_logs job_id={job_id} application_id={job.application_id} " + f"chars={len(tailed)}" + ) + return tailed diff --git a/tests/unit/test_logs_tool.py b/tests/unit/test_logs_tool.py new file mode 100644 index 0000000..3c349de --- /dev/null +++ b/tests/unit/test_logs_tool.py @@ -0,0 +1,47 @@ +# coding=utf-8 +from datetime import datetime +from unittest.mock import patch + +from spark_executor.core.job_store import JobStore +from spark_executor.models import Job +from spark_executor.tools import logs + + +def _seed(job_id="abc", app_id="application_1"): + logs.store = JobStore() + logs.store.put( + Job( + job_id=job_id, + application_id=app_id, + script_path="/tmp/j.py", + queue="default", + submit_time=datetime(2026, 6, 24), + connection="prod", + ) + ) + + +def test_get_job_logs_tails_to_default_5000(): + _seed() + big = "x" * 8000 + "\nEND" + with patch("spark_executor.tools.logs.get_application_logs", return_value=big): + out = logs.get_job_logs("abc") + assert out.endswith("END") + assert len(out) == 5000 + + +def test_get_job_logs_respects_custom_tail_chars(): + _seed() + with patch( + "spark_executor.tools.logs.get_application_logs", + return_value="0123456789", + ): + out = logs.get_job_logs("abc", tail_chars=3) + assert out == "789" + + +def test_get_job_logs_raises_for_unknown_job(): + _seed() + import pytest + with pytest.raises(KeyError): + logs.get_job_logs("missing") From 0d60f9b246a3c7700b1e4fe2d91619a09a93d69b Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:44:10 +0800 Subject: [PATCH 18/23] feat: add kill_job tool --- spark_executor/tools/kill.py | 23 +++++++++++++++++++++++ tests/unit/test_kill_tool.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+) create mode 100644 spark_executor/tools/kill.py create mode 100644 tests/unit/test_kill_tool.py diff --git a/spark_executor/tools/kill.py b/spark_executor/tools/kill.py new file mode 100644 index 0000000..da692c1 --- /dev/null +++ b/spark_executor/tools/kill.py @@ -0,0 +1,23 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" +from common.logging import logger +from spark_executor.core.job_store import JobStore +from spark_executor.core.yarn_client import kill_application + +store = JobStore() + + +def kill_job(job_id: str) -> dict[str, str]: + job = store.get(job_id) + if job is None: + raise KeyError(f"Unknown job_id: {job_id}") + kill_application(job.application_id) + logger.info(f"kill_job job_id={job_id} application_id={job.application_id}") + return { + "job_id": job_id, + "application_id": job.application_id, + "status": "KILLED", + } diff --git a/tests/unit/test_kill_tool.py b/tests/unit/test_kill_tool.py new file mode 100644 index 0000000..bf7e490 --- /dev/null +++ b/tests/unit/test_kill_tool.py @@ -0,0 +1,36 @@ +# coding=utf-8 +from datetime import datetime +from unittest.mock import patch + +from spark_executor.core.job_store import JobStore +from spark_executor.models import Job +from spark_executor.tools import kill + + +def test_kill_job_calls_yarn_kill(): + kill.store = JobStore() + kill.store.put( + Job( + job_id="abc", + application_id="application_1", + script_path="/tmp/j.py", + queue="default", + submit_time=datetime(2026, 6, 24), + connection="prod", + ) + ) + with patch("spark_executor.tools.kill.kill_application") as m: + result = kill.kill_job("abc") + m.assert_called_once_with("application_1") + assert result == { + "job_id": "abc", + "application_id": "application_1", + "status": "KILLED", + } + + +def test_kill_job_raises_for_unknown_job(): + kill.store = JobStore() + import pytest + with pytest.raises(KeyError): + kill.kill_job("missing") From cc7eafaa16912b3952c0a8a7a41ff600510f3c5e Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:50:51 +0800 Subject: [PATCH 19/23] feat: register 12 MCP tools on FastAPI app --- spark_executor/server.py | 107 ++++++++++++++++++++++++++- tests/integration/__init__.py | 5 ++ tests/integration/test_mcp_routes.py | 33 +++++++++ 3 files changed, 143 insertions(+), 2 deletions(-) create mode 100644 tests/integration/__init__.py create mode 100644 tests/integration/test_mcp_routes.py diff --git a/spark_executor/server.py b/spark_executor/server.py index 6b69e74..fe18745 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -1,12 +1,115 @@ # coding=utf-8 """ @Time :2026/6/24 -@Author :tao.chen +@Author :tao.chen """ from fastapi import FastAPI +from spark_executor.tools.connections import ( + delete_connection, + get_connection, + list_connections, + save_connection, +) +from spark_executor.tools.kill import kill_job +from spark_executor.tools.logs import get_job_logs +from spark_executor.tools.status import get_job_status +from spark_executor.tools.submit import ( + cancel_pending_job, + confirm_submit_job, + get_pending_job, + list_pending_jobs, + prepare_submit_job, +) + app = FastAPI(title="Spark Executor", version="0.0.1", description="Spark Executor MCP Server") + @app.get("/health") def health_check(): - return {"status": "ok"} \ No newline at end of file + return {"status": "ok"} + + +# MCP tool routes. fastapi-mcp discovers these and registers them as MCP tools. + +# --- Pending submission flow (two-step submit) --- + +@app.post("/prepare_submit_job") +def _prepare_submit_job(connection: str, script_path: str, queue: str = "default", + executor_memory: str = "4G", executor_cores: int = 2, + num_executors: int = 2): + return prepare_submit_job( + connection=connection, + script_path=script_path, + queue=queue, + executor_memory=executor_memory, + executor_cores=executor_cores, + num_executors=num_executors, + ) + + +@app.post("/confirm_submit_job") +def _confirm_submit_job(pending_id: str): + return confirm_submit_job(pending_id=pending_id) + + +@app.post("/list_pending_jobs") +def _list_pending_jobs(): + return list_pending_jobs() + + +@app.post("/get_pending_job") +def _get_pending_job(pending_id: str): + return get_pending_job(pending_id) + + +@app.post("/cancel_pending_job") +def _cancel_pending_job(pending_id: str): + return cancel_pending_job(pending_id) + + +# --- Spark job tools --- + +@app.post("/get_job_status") +def _get_job_status(job_id: str): + return get_job_status(job_id) + + +@app.post("/get_job_logs") +def _get_job_logs(job_id: str, tail_chars: int = 5000): + return get_job_logs(job_id, tail_chars=tail_chars) + + +@app.post("/kill_job") +def _kill_job(job_id: str): + return kill_job(job_id) + + +# --- Connection management tools --- + +@app.post("/save_connection") +def _save_connection(name: str, master: str, deploy_mode: str = "cluster", + yarn_rm_url: str | None = None, + spark_conf: dict[str, str] | None = None): + return save_connection( + name=name, + master=master, + deploy_mode=deploy_mode, + yarn_rm_url=yarn_rm_url, + spark_conf=spark_conf, + ) + + +@app.post("/list_connections") +def _list_connections(): + return list_connections() + + +@app.post("/get_connection") +def _get_connection(name: str): + return get_connection(name) + + +@app.post("/delete_connection") +def _delete_connection(name: str): + return delete_connection(name) diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 0000000..36f7e1b --- /dev/null +++ b/tests/integration/__init__.py @@ -0,0 +1,5 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen +""" diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py new file mode 100644 index 0000000..dd1ded9 --- /dev/null +++ b/tests/integration/test_mcp_routes.py @@ -0,0 +1,33 @@ +# coding=utf-8 +from fastapi.testclient import TestClient + +from spark_executor.server import app + + +def test_health_still_present(): + c = TestClient(app) + r = c.get("/health") + assert r.status_code == 200 + assert r.json() == {"status": "ok"} + + +def test_twelve_tool_routes_registered(): + paths = {r.path for r in app.routes} + for path in ( + # pending-submission flow (5) + "/prepare_submit_job", + "/confirm_submit_job", + "/list_pending_jobs", + "/get_pending_job", + "/cancel_pending_job", + # job lifecycle (3) + "/get_job_status", + "/get_job_logs", + "/kill_job", + # connection management (4) + "/save_connection", + "/list_connections", + "/get_connection", + "/delete_connection", + ): + assert path in paths, f"missing MCP tool route: {path}" From 6b9f2782944504c7184d2246354eecdb3d8e3379 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:56:31 +0800 Subject: [PATCH 20/23] fix: switch FastAPI routes to Pydantic body models for MCP tools/call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The plan's route signatures used query parameters, which worked for direct HTTP callers and unit tests, but fastapi-mcp's HTTP transport passes tools/call arguments as a JSON body. dict-typed parameters like spark_conf arrived as a string and the route returned 422. Refactor each route to take a single Pydantic body model (saved in spark_executor/tools/requests.py). Underlying tool functions unchanged. Integration tests in tests/integration/test_mcp_routes.py now exercise the full body-based roundtrip (save_connection with spark_conf, prepare → list → get, list_connections with empty body). --- spark_executor/server.py | 74 ++++++++++++-------------- spark_executor/tools/requests.py | 50 ++++++++++++++++++ tests/integration/test_mcp_routes.py | 79 ++++++++++++++++++++++++++++ 3 files changed, 164 insertions(+), 39 deletions(-) create mode 100644 spark_executor/tools/requests.py diff --git a/spark_executor/server.py b/spark_executor/server.py index fe18745..578c733 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -13,6 +13,15 @@ from spark_executor.tools.connections import ( ) from spark_executor.tools.kill import kill_job from spark_executor.tools.logs import get_job_logs +from spark_executor.tools.requests import ( + ConnectionNameRequest, + EmptyRequest, + GetJobLogsRequest, + JobIdRequest, + PendingIdRequest, + PrepareSubmitJobRequest, + SaveConnectionRequest, +) from spark_executor.tools.status import get_job_status from spark_executor.tools.submit import ( cancel_pending_job, @@ -31,85 +40,72 @@ def health_check(): # MCP tool routes. fastapi-mcp discovers these and registers them as MCP tools. +# Each route takes a single Pydantic body model so tools/call (which sends args +# as JSON body) works for every tool, including those with dict-typed params +# like spark_conf. # --- Pending submission flow (two-step submit) --- @app.post("/prepare_submit_job") -def _prepare_submit_job(connection: str, script_path: str, queue: str = "default", - executor_memory: str = "4G", executor_cores: int = 2, - num_executors: int = 2): - return prepare_submit_job( - connection=connection, - script_path=script_path, - queue=queue, - executor_memory=executor_memory, - executor_cores=executor_cores, - num_executors=num_executors, - ) +def _prepare_submit_job(req: PrepareSubmitJobRequest): + return prepare_submit_job(**req.model_dump()) @app.post("/confirm_submit_job") -def _confirm_submit_job(pending_id: str): - return confirm_submit_job(pending_id=pending_id) +def _confirm_submit_job(req: PendingIdRequest): + return confirm_submit_job(pending_id=req.pending_id) @app.post("/list_pending_jobs") -def _list_pending_jobs(): +def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()): return list_pending_jobs() @app.post("/get_pending_job") -def _get_pending_job(pending_id: str): - return get_pending_job(pending_id) +def _get_pending_job(req: PendingIdRequest): + return get_pending_job(req.pending_id) @app.post("/cancel_pending_job") -def _cancel_pending_job(pending_id: str): - return cancel_pending_job(pending_id) +def _cancel_pending_job(req: PendingIdRequest): + return cancel_pending_job(req.pending_id) # --- Spark job tools --- @app.post("/get_job_status") -def _get_job_status(job_id: str): - return get_job_status(job_id) +def _get_job_status(req: JobIdRequest): + return get_job_status(req.job_id) @app.post("/get_job_logs") -def _get_job_logs(job_id: str, tail_chars: int = 5000): - return get_job_logs(job_id, tail_chars=tail_chars) +def _get_job_logs(req: GetJobLogsRequest): + return get_job_logs(req.job_id, tail_chars=req.tail_chars) @app.post("/kill_job") -def _kill_job(job_id: str): - return kill_job(job_id) +def _kill_job(req: JobIdRequest): + return kill_job(req.job_id) # --- Connection management tools --- @app.post("/save_connection") -def _save_connection(name: str, master: str, deploy_mode: str = "cluster", - yarn_rm_url: str | None = None, - spark_conf: dict[str, str] | None = None): - return save_connection( - name=name, - master=master, - deploy_mode=deploy_mode, - yarn_rm_url=yarn_rm_url, - spark_conf=spark_conf, - ) +def _save_connection(req: SaveConnectionRequest): + # exclude_none so we don't overwrite the function's default with explicit None + return save_connection(**req.model_dump(exclude_none=True)) @app.post("/list_connections") -def _list_connections(): +def _list_connections(_req: EmptyRequest = EmptyRequest()): return list_connections() @app.post("/get_connection") -def _get_connection(name: str): - return get_connection(name) +def _get_connection(req: ConnectionNameRequest): + return get_connection(req.name) @app.post("/delete_connection") -def _delete_connection(name: str): - return delete_connection(name) +def _delete_connection(req: ConnectionNameRequest): + return delete_connection(req.name) diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py new file mode 100644 index 0000000..d7c0d55 --- /dev/null +++ b/spark_executor/tools/requests.py @@ -0,0 +1,50 @@ +# coding=utf-8 +""" +@Time :2026/6/24 +@Author :tao.chen + +Pydantic request models for the FastAPI route layer. The underlying tool +functions in tools/*.py still take keyword arguments; these models exist only +so fastapi-mcp can call the routes via tools/call (which sends args as a +JSON body) without 422-ing on dict-typed parameters like spark_conf. +""" +from pydantic import BaseModel, Field + + +class EmptyRequest(BaseModel): + """Used for tools that take no arguments (list_connections, list_pending_jobs).""" + pass + + +class SaveConnectionRequest(BaseModel): + name: str + master: str + deploy_mode: str = "cluster" + yarn_rm_url: str | None = None + spark_conf: dict[str, str] | None = None + + +class PrepareSubmitJobRequest(BaseModel): + connection: str + script_path: str + queue: str = "default" + executor_memory: str = "4G" + executor_cores: int = 2 + num_executors: int = 2 + + +class PendingIdRequest(BaseModel): + pending_id: str + + +class JobIdRequest(BaseModel): + job_id: str + + +class GetJobLogsRequest(BaseModel): + job_id: str + tail_chars: int = 5000 + + +class ConnectionNameRequest(BaseModel): + name: str diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index dd1ded9..4b85cf8 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -1,7 +1,28 @@ # coding=utf-8 +from pathlib import Path + +import pytest from fastapi.testclient import TestClient +from spark_executor.core import connection_store, pending_store +from spark_executor.core.connection_store import ConnectionStore +from spark_executor.core.pending_store import PendingStore from spark_executor.server import app +from spark_executor.tools import connections, submit + + +@pytest.fixture(autouse=True) +def _fresh_data(tmp_path: Path, monkeypatch): + """Reset both stores to a fresh tmp dir and rebind the singletons that + the tool modules captured at import time.""" + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", ConnectionStore()) + monkeypatch.setattr(pending_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(pending_store, "store", PendingStore()) + # Rebind imports inside the tool modules (captured at import time) + connections.store = connection_store.store + submit.conn_store = connection_store.store + submit.pending_store = pending_store.store def test_health_still_present(): @@ -31,3 +52,61 @@ def test_twelve_tool_routes_registered(): "/delete_connection", ): assert path in paths, f"missing MCP tool route: {path}" + + +# --- End-to-end body-based calls (the gap the route-registration test missed) --- + +def test_save_connection_accepts_dict_spark_conf_in_body(): + """The original query-param signature returned 422 for spark_conf dicts; + body models make tools/call roundtrip cleanly.""" + c = TestClient(app) + r = c.post( + "/save_connection", + json={ + "name": "prod", + "master": "yarn", + "deploy_mode": "cluster", + "spark_conf": {"spark.sql.shuffle.partitions": "200"}, + }, + ) + assert r.status_code == 200, r.text + assert r.json() == {"name": "prod", "status": "SAVED"} + + +def test_prepare_submit_job_works_via_body(): + c = TestClient(app) + c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + r = c.post( + "/prepare_submit_job", + json={"connection": "prod", "script_path": "/tmp/demo.py", "queue": "research"}, + ) + assert r.status_code == 200, r.text + body = r.json() + assert body["status"] == "PENDING" + assert body["pending_id"].startswith("p_") + assert body["parameters"]["queue"] == "research" + assert body["parameters"]["master"] == "yarn" # snapshotted from connection + + +def test_list_and_get_pending_job_roundtrip_via_body(): + c = TestClient(app) + c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + prep = c.post( + "/prepare_submit_job", + json={"connection": "prod", "script_path": "/tmp/a.py"}, + ).json() + pid = prep["pending_id"] + + listed = c.post("/list_pending_jobs", json={}).json() + assert any(p["pending_id"] == pid for p in listed) + + got = c.post("/get_pending_job", json={"pending_id": pid}).json() + assert got["script_path"] == "/tmp/a.py" + assert got["status"] == "PENDING" + + +def test_list_connections_works_with_empty_body(): + c = TestClient(app) + r = c.post("/list_connections", json={}) + assert r.status_code == 200 + assert r.json() == [] From fad296591c8be3f036d0e9f8d7c1b870ba164a7c Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 14:59:00 +0800 Subject: [PATCH 21/23] fix: map KeyError->404 and ValueError->400 for MCP clients MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tool functions raise KeyError for unknown ids and ValueError for invalid state transitions. Without handlers, FastAPI returned a bare 500 "Internal Server Error" to MCP clients — useless for an agent that needs to know whether the id was wrong vs the cluster was down. KeyError -> 404 with the message ("Unknown job_id: foo"). ValueError -> 400 with the message (e.g. "pending_id p_xyz is in status 'CANCELLED', not PENDING"). --- spark_executor/server.py | 20 ++++++++- tests/integration/test_mcp_routes.py | 66 ++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+), 1 deletion(-) diff --git a/spark_executor/server.py b/spark_executor/server.py index 578c733..16ad11d 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -3,7 +3,8 @@ @Time :2026/6/24 @Author :tao.chen """ -from fastapi import FastAPI +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse from spark_executor.tools.connections import ( delete_connection, @@ -34,6 +35,23 @@ from spark_executor.tools.submit import ( app = FastAPI(title="Spark Executor", version="0.0.1", description="Spark Executor MCP Server") +# --- Exception handlers: translate tool-layer errors into proper HTTP statuses --- +# +# Tool functions raise KeyError for "unknown id" (job_id, pending_id, connection +# name) and ValueError for invalid state transitions (e.g. confirming a +# CANCELLED pending). Without these handlers FastAPI would map them to a bare +# 500 "Internal Server Error" which is useless to MCP clients. + +@app.exception_handler(KeyError) +async def _keyerror_handler(_request: Request, exc: KeyError) -> JSONResponse: + return JSONResponse(status_code=404, content={"detail": str(exc)}) + + +@app.exception_handler(ValueError) +async def _valueerror_handler(_request: Request, exc: ValueError) -> JSONResponse: + return JSONResponse(status_code=400, content={"detail": str(exc)}) + + @app.get("/health") def health_check(): return {"status": "ok"} diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index 4b85cf8..c77a975 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -110,3 +110,69 @@ def test_list_connections_works_with_empty_body(): r = c.post("/list_connections", json={}) assert r.status_code == 200 assert r.json() == [] + + +# --- Exception handlers: KeyError -> 404, ValueError -> 400 --- + +def test_unknown_job_id_returns_404(): + c = TestClient(app) + r = c.post("/get_job_status", json={"job_id": "missing"}) + assert r.status_code == 404 + assert "missing" in r.json()["detail"] + + +def test_unknown_pending_id_returns_404(): + c = TestClient(app) + r = c.post("/get_pending_job", json={"pending_id": "p_nope"}) + assert r.status_code == 404 + assert "p_nope" in r.json()["detail"] + + +def test_unknown_connection_name_returns_404(): + c = TestClient(app) + r = c.post("/get_connection", json={"name": "nope"}) + assert r.status_code == 404 + assert "nope" in r.json()["detail"] + + +def test_unknown_connection_in_prepare_returns_404(): + c = TestClient(app) + r = c.post( + "/prepare_submit_job", + json={"connection": "nope", "script_path": "/tmp/x.py"}, + ) + assert r.status_code == 404 + assert "nope" in r.json()["detail"] + + +def test_confirm_non_pending_returns_400(): + """A CANCELLED pending should refuse confirm; FastAPI should surface ValueError as 400.""" + c = TestClient(app) + c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + prep = c.post( + "/prepare_submit_job", + json={"connection": "prod", "script_path": "/tmp/x.py"}, + ).json() + pid = prep["pending_id"] + c.post("/cancel_pending_job", json={"pending_id": pid}) + r = c.post("/confirm_submit_job", json={"pending_id": pid}) + assert r.status_code == 400 + assert "CANCELLED" in r.json()["detail"] + + +def test_cancel_already_submitted_returns_400(): + """Cancelling a SUBMITTED pending should refuse with ValueError -> 400.""" + c = TestClient(app) + c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + prep = c.post( + "/prepare_submit_job", + json={"connection": "prod", "script_path": "/tmp/x.py"}, + ).json() + pid = prep["pending_id"] + # Simulate a SUBMITTED state by mutating the pending directly + p = pending_store.store.get(pid) + p.status = "SUBMITTED" + pending_store.store.save(p) + r = c.post("/cancel_pending_job", json={"pending_id": pid}) + assert r.status_code == 400 + assert "SUBMITTED" in r.json()["detail"] From 5c7dcf57cce1c5a69f7c49786ef1fa5277d69842 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 15:02:54 +0800 Subject: [PATCH 22/23] feat: structured DEBUG/INFO logging via loguru - common/logging.py: improve format to timestamp|LEVEL|module:func:line - message - core/ layer: DEBUG log every subprocess invocation (cmd, rc, byte counts), JSON load/dump events, parsed application_id. ERROR log on failures. - tools/ layer: DEBUG log every public tool entry with key parameters, INFO log on business outcomes (saved/submitted/killed/...). - New tests/unit/test_logging.py: capture loguru output via in-memory sink and assert DEBUG + INFO messages are emitted for representative flows. --- common/logging.py | 23 ++++++++- spark_executor/core/connection_store.py | 8 +++- spark_executor/core/log_parser.py | 8 ++++ spark_executor/core/pending_store.py | 11 ++++- spark_executor/core/spark_submit.py | 9 ++++ spark_executor/core/yarn_client.py | 27 ++++++++++- spark_executor/tools/connections.py | 9 +++- spark_executor/tools/kill.py | 3 +- spark_executor/tools/logs.py | 5 +- spark_executor/tools/status.py | 3 +- spark_executor/tools/submit.py | 28 ++++++++--- tests/unit/test_logging.py | 63 +++++++++++++++++++++++++ 12 files changed, 178 insertions(+), 19 deletions(-) create mode 100644 tests/unit/test_logging.py diff --git a/common/logging.py b/common/logging.py index 1325ddf..50bc3f2 100644 --- a/common/logging.py +++ b/common/logging.py @@ -1,10 +1,29 @@ # coding=utf-8 """ @Time :2026/6/24 -@Author :tao.chen +@Author :tao.chen + +Process-wide loguru configuration. Import `logger` from here in every +module instead of instantiating new loggers. + +Levels used in this project: + DEBUG - entry/exit of public tools, subprocess commands, file I/O paths + INFO - business events (job submitted, status changed, connection saved) + WARNING - recoverable problems (transient YARN issues, retry-able) + ERROR - raised exceptions (caller will see the traceback) """ import sys + from loguru import logger logger.remove() -logger.add(sys.stderr, level="DEBUG") \ No newline at end of file +logger.add( + sys.stderr, + level="DEBUG", + format=( + "{time:HH:mm:ss.SSS} | " + "{level: <7} | " + "{name}:{function}:{line} - " + "{message}" + ), +) diff --git a/spark_executor/core/connection_store.py b/spark_executor/core/connection_store.py index 503f67b..15729a5 100644 --- a/spark_executor/core/connection_store.py +++ b/spark_executor/core/connection_store.py @@ -9,6 +9,7 @@ import tempfile from pathlib import Path from threading import Lock +from common.logging import logger from spark_executor.models import Connection DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") @@ -37,8 +38,10 @@ class ConnectionStore: def _load(self) -> dict[str, Connection]: if not self.path.exists(): + logger.debug(f"ConnectionStore._load: file {self.path} absent, returning empty") return {} raw = json.loads(self.path.read_text(encoding="utf-8")) + logger.debug(f"ConnectionStore._load: loaded {len(raw)} records from {self.path}") return {name: Connection.model_validate(c) for name, c in raw.items()} def _dump(self, records: dict[str, Connection]) -> None: @@ -52,6 +55,7 @@ class ConnectionStore: with os.fdopen(fd, "w", encoding="utf-8") as f: json.dump(payload, f, indent=2, ensure_ascii=False) os.replace(tmp_path, self.path) + logger.debug(f"ConnectionStore._dump: wrote {len(records)} records to {self.path}") except Exception: if os.path.exists(tmp_path): os.unlink(tmp_path) @@ -70,6 +74,7 @@ class ConnectionStore: records = self._load() records[conn.name] = conn self._dump(records) + logger.info(f"connection saved name={conn.name} master={conn.master}") def delete(self, name: str) -> bool: with self._lock: @@ -78,7 +83,8 @@ class ConnectionStore: return False del records[name] self._dump(records) - return True + logger.info(f"connection deleted name={name}") + return True # Module-level singleton; replaced in tests. diff --git a/spark_executor/core/log_parser.py b/spark_executor/core/log_parser.py index ad1b72e..306966b 100644 --- a/spark_executor/core/log_parser.py +++ b/spark_executor/core/log_parser.py @@ -5,6 +5,8 @@ """ import re +from common.logging import logger + _APP_ID_RE = re.compile(r"Submitted application (\S+)") _TRACKING_URL_RE = re.compile(r"tracking URL:\s+(\S+)") @@ -20,14 +22,20 @@ def parse_spark_submit_output(stderr: str) -> tuple[str, str | None]: else: url_match = _TRACKING_URL_RE.search(stderr) if not url_match: + logger.error("Could not find application_id in spark-submit output") raise ValueError("Could not find application_id in spark-submit output") # tracking URL is of the form http://rm:8088/proxy/application_xxx/ tracking_url = url_match.group(1) tail = tracking_url.rstrip("/").rsplit("/", 1)[-1] if not tail.startswith("application_"): + logger.error("Could not find application_id in spark-submit output (bad tracking URL)") raise ValueError("Could not find application_id in spark-submit output") application_id = tail url_match = _TRACKING_URL_RE.search(stderr) tracking_url = url_match.group(1) if url_match else None + logger.debug( + f"parse_spark_submit_output -> application_id={application_id} " + f"tracking_url={tracking_url}" + ) return application_id, tracking_url diff --git a/spark_executor/core/pending_store.py b/spark_executor/core/pending_store.py index 661229e..8487881 100644 --- a/spark_executor/core/pending_store.py +++ b/spark_executor/core/pending_store.py @@ -9,6 +9,7 @@ import tempfile from pathlib import Path from threading import Lock +from common.logging import logger from spark_executor.models import PendingSubmission DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data") @@ -33,8 +34,10 @@ class PendingStore: def _load(self) -> dict[str, PendingSubmission]: if not self.path.exists(): + logger.debug(f"PendingStore._load: file {self.path} absent, returning empty") return {} raw = json.loads(self.path.read_text(encoding="utf-8")) + logger.debug(f"PendingStore._load: loaded {len(raw)} records from {self.path}") return {pid: PendingSubmission.model_validate(p) for pid, p in raw.items()} def _dump(self, records: dict[str, PendingSubmission]) -> None: @@ -47,6 +50,7 @@ class PendingStore: with os.fdopen(fd, "w", encoding="utf-8") as f: json.dump(payload, f, indent=2, ensure_ascii=False, default=str) os.replace(tmp_path, self.path) + logger.debug(f"PendingStore._dump: wrote {len(records)} records to {self.path}") except Exception: if os.path.exists(tmp_path): os.unlink(tmp_path) @@ -65,6 +69,10 @@ class PendingStore: records = self._load() records[pending.pending_id] = pending self._dump(records) + logger.info( + f"pending saved pending_id={pending.pending_id} " + f"status={pending.status} connection={pending.connection}" + ) def delete(self, pending_id: str) -> bool: with self._lock: @@ -73,7 +81,8 @@ class PendingStore: return False del records[pending_id] self._dump(records) - return True + logger.info(f"pending deleted pending_id={pending_id}") + return True # Module-level singleton; replaced in tests. diff --git a/spark_executor/core/spark_submit.py b/spark_executor/core/spark_submit.py index 559845d..d0f759d 100644 --- a/spark_executor/core/spark_submit.py +++ b/spark_executor/core/spark_submit.py @@ -5,6 +5,8 @@ """ import subprocess +from common.logging import logger + class SparkSubmitError(Exception): """Raised when spark-submit exits with a non-zero return code.""" @@ -33,12 +35,19 @@ def build_spark_submit_command( for key, value in (spark_conf or {}).items(): cmd.extend(["--conf", f"{key}={value}"]) cmd.append(script_path) + logger.debug(f"build_spark_submit_command -> {cmd}") return cmd def run_spark_submit(cmd: list[str]) -> "subprocess.CompletedProcess[str]": + logger.debug(f"run_spark_submit exec: {cmd}") result = subprocess.run(cmd, capture_output=True, text=True, errors="replace") + logger.debug( + f"run_spark_submit done rc={result.returncode} " + f"stdout_len={len(result.stdout)} stderr_len={len(result.stderr)}" + ) if result.returncode != 0: + logger.error(f"spark-submit failed (rc={result.returncode}): {result.stderr[:500]}") raise SparkSubmitError( f"spark-submit failed (rc={result.returncode}): {result.stderr}" ) diff --git a/spark_executor/core/yarn_client.py b/spark_executor/core/yarn_client.py index 15bcec1..89ab071 100644 --- a/spark_executor/core/yarn_client.py +++ b/spark_executor/core/yarn_client.py @@ -6,6 +6,8 @@ import re import subprocess +from common.logging import logger + class YarnError(Exception): """Raised when a yarn CLI invocation fails.""" @@ -15,33 +17,54 @@ _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") + logger.debug(f"yarn _run exec: {cmd}") + result = subprocess.run(cmd, capture_output=True, text=True, errors="replace") + logger.debug( + f"yarn _run done rc={result.returncode} " + f"stdout_len={len(result.stdout)} stderr_len={len(result.stderr)}" + ) + return result def get_application_status(application_id: str) -> tuple[str, str]: proc = _run(["yarn", "application", "-status", application_id]) if proc.returncode != 0: + logger.error( + f"yarn application -status failed (rc={proc.returncode}) for " + f"{application_id}: {proc.stderr[:500]}" + ) 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 + state = match.group(1) + logger.info(f"yarn status {application_id} -> {state}") + return state, proc.stdout def get_application_logs(application_id: str) -> str: proc = _run(["yarn", "logs", "-applicationId", application_id]) if proc.returncode != 0: + logger.error( + f"yarn logs failed (rc={proc.returncode}) for {application_id}: {proc.stderr[:500]}" + ) raise YarnError( f"yarn logs failed (rc={proc.returncode}): {proc.stderr}" ) + logger.info(f"yarn logs {application_id} -> {len(proc.stdout)} chars") return proc.stdout def kill_application(application_id: str) -> None: proc = _run(["yarn", "application", "-kill", application_id]) if proc.returncode != 0: + logger.error( + f"yarn application -kill failed (rc={proc.returncode}) for " + f"{application_id}: {proc.stderr[:500]}" + ) raise YarnError( f"yarn application -kill failed (rc={proc.returncode}): {proc.stderr}" ) + logger.info(f"yarn kill {application_id} -> ok") diff --git a/spark_executor/tools/connections.py b/spark_executor/tools/connections.py index ccff910..f76c8f9 100644 --- a/spark_executor/tools/connections.py +++ b/spark_executor/tools/connections.py @@ -16,6 +16,10 @@ def save_connection( yarn_rm_url: str | None = None, spark_conf: dict[str, str] | None = None, ) -> dict[str, str]: + logger.debug( + f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} " + f"yarn_rm_url={yarn_rm_url} spark_conf_keys={list((spark_conf or {}).keys())}" + ) conn = Connection( name=name, master=master, @@ -24,15 +28,16 @@ def save_connection( spark_conf=spark_conf or {}, ) store.save(conn) - logger.info(f"save_connection name={name} master={master}") return {"name": name, "status": "SAVED"} def list_connections() -> list[dict[str, object]]: + logger.debug("list_connections enter") return [c.model_dump() for c in store.list_all()] def get_connection(name: str) -> dict[str, object]: + logger.debug(f"get_connection enter name={name}") conn = store.get(name) if conn is None: raise KeyError(f"Unknown connection: {name}") @@ -40,8 +45,8 @@ def get_connection(name: str) -> dict[str, object]: def delete_connection(name: str) -> dict[str, str]: + logger.debug(f"delete_connection enter name={name}") removed = store.delete(name) if not removed: raise KeyError(f"Unknown connection: {name}") - logger.info(f"delete_connection name={name}") return {"name": name, "status": "DELETED"} diff --git a/spark_executor/tools/kill.py b/spark_executor/tools/kill.py index da692c1..8d9409e 100644 --- a/spark_executor/tools/kill.py +++ b/spark_executor/tools/kill.py @@ -11,11 +11,12 @@ store = JobStore() def kill_job(job_id: str) -> dict[str, str]: + logger.debug(f"kill_job enter job_id={job_id}") job = store.get(job_id) if job is None: raise KeyError(f"Unknown job_id: {job_id}") kill_application(job.application_id) - logger.info(f"kill_job job_id={job_id} application_id={job.application_id}") + logger.info(f"kill_job ok job_id={job_id} application_id={job.application_id}") return { "job_id": job_id, "application_id": job.application_id, diff --git a/spark_executor/tools/logs.py b/spark_executor/tools/logs.py index ba0f8e0..2834878 100644 --- a/spark_executor/tools/logs.py +++ b/spark_executor/tools/logs.py @@ -11,13 +11,14 @@ store = JobStore() def get_job_logs(job_id: str, tail_chars: int = 5000) -> str: + logger.debug(f"get_job_logs enter job_id={job_id} tail_chars={tail_chars}") job = store.get(job_id) if job is None: raise KeyError(f"Unknown job_id: {job_id}") full = get_application_logs(job.application_id) tailed = full[-tail_chars:] if len(full) > tail_chars else full logger.info( - f"get_job_logs job_id={job_id} application_id={job.application_id} " - f"chars={len(tailed)}" + f"get_job_logs ok job_id={job_id} application_id={job.application_id} " + f"full_chars={len(full)} returned_chars={len(tailed)}" ) return tailed diff --git a/spark_executor/tools/status.py b/spark_executor/tools/status.py index 048ff1a..17b2f86 100644 --- a/spark_executor/tools/status.py +++ b/spark_executor/tools/status.py @@ -12,9 +12,10 @@ store = JobStore() def get_job_status(job_id: str) -> JobStatus: + logger.debug(f"get_job_status enter job_id={job_id}") job = store.get(job_id) if job is None: raise KeyError(f"Unknown job_id: {job_id}") state, raw = get_application_status(job.application_id) - logger.info(f"get_job_status job_id={job_id} state={state}") + logger.info(f"get_job_status ok job_id={job_id} application_id={job.application_id} state={state}") return JobStatus(application_id=job.application_id, state=state, raw=raw) diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py index 8e09d12..db72311 100644 --- a/spark_executor/tools/submit.py +++ b/spark_executor/tools/submit.py @@ -4,6 +4,7 @@ @Author :tao.chen """ import secrets +import uuid from datetime import datetime from common.logging import logger @@ -17,8 +18,6 @@ from spark_executor.core.spark_submit import ( run_spark_submit, ) from spark_executor.models import Job, PendingSubmission, SubmitResult -import uuid -from datetime import datetime as _datetime def _new_pending_id() -> str: @@ -35,6 +34,11 @@ def prepare_submit_job( num_executors: int = 2, ) -> dict[str, object]: """Snapshot connection params and persist a PendingSubmission. Does NOT submit.""" + logger.debug( + f"prepare_submit_job enter connection={connection} script_path={script_path} " + f"queue={queue} executor_memory={executor_memory} executor_cores={executor_cores} " + f"num_executors={num_executors}" + ) conn = conn_store.get(connection) if conn is None: raise KeyError(f"Unknown connection: {connection}") @@ -56,7 +60,8 @@ def prepare_submit_job( ) pending_store.save(pending) logger.info( - f"prepare_submit_job pending_id={pending_id} connection={connection} master={conn.master}" + f"prepare_submit_job ok pending_id={pending_id} connection={connection} " + f"master={conn.master} script_path={script_path}" ) return { "pending_id": pending_id, @@ -71,6 +76,7 @@ job_store: JobStore = JobStore() def confirm_submit_job(*, pending_id: str) -> SubmitResult: """Actually invoke spark-submit for a previously-prepared PendingSubmission.""" + logger.debug(f"confirm_submit_job enter pending_id={pending_id}") pending = pending_store.get(pending_id) if pending is None: raise KeyError(f"Unknown pending_id: {pending_id}") @@ -89,13 +95,17 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult: num_executors=pending.num_executors, spark_conf=pending.spark_conf, ) - logger.info(f"confirm_submit_job pending_id={pending_id} cmd={cmd}") + logger.info( + f"confirm_submit_job start pending_id={pending_id} " + f"application_target={pending.master} script_path={pending.script_path}" + ) try: result = run_spark_submit(cmd) except SparkSubmitError as exc: pending.status = "FAILED" pending.error = str(exc) pending_store.save(pending) + logger.error(f"confirm_submit_job failed pending_id={pending_id} err={exc}") raise application_id, tracking_url = parse_spark_submit_output(result.stderr) @@ -107,7 +117,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult: application_id=application_id, script_path=pending.script_path, queue=pending.queue, - submit_time=_datetime.utcnow(), + submit_time=datetime.utcnow(), connection=pending.connection, ) ) @@ -117,7 +127,8 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult: pending.application_id = application_id pending_store.save(pending) logger.info( - f"confirm_submit_job pending_id={pending_id} job_id={job_id} application_id={application_id}" + f"confirm_submit_job ok pending_id={pending_id} job_id={job_id} " + f"application_id={application_id}" ) return SubmitResult( job_id=job_id, @@ -127,10 +138,12 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult: def list_pending_jobs() -> list[dict[str, object]]: + logger.debug("list_pending_jobs enter") return [p.model_dump() for p in pending_store.list_all()] def get_pending_job(pending_id: str) -> dict[str, object]: + logger.debug(f"get_pending_job enter pending_id={pending_id}") p = pending_store.get(pending_id) if p is None: raise KeyError(f"Unknown pending_id: {pending_id}") @@ -138,6 +151,7 @@ def get_pending_job(pending_id: str) -> dict[str, object]: def cancel_pending_job(pending_id: str) -> dict[str, str]: + logger.debug(f"cancel_pending_job enter pending_id={pending_id}") p = pending_store.get(pending_id) if p is None: raise KeyError(f"Unknown pending_id: {pending_id}") @@ -147,5 +161,5 @@ def cancel_pending_job(pending_id: str) -> dict[str, str]: ) p.status = "CANCELLED" pending_store.save(p) - logger.info(f"cancel_pending_job pending_id={pending_id}") + logger.info(f"cancel_pending_job ok pending_id={pending_id}") return {"pending_id": pending_id, "status": "CANCELLED"} diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py new file mode 100644 index 0000000..44b746b --- /dev/null +++ b/tests/unit/test_logging.py @@ -0,0 +1,63 @@ +# coding=utf-8 +import io +from pathlib import Path + +import pytest + +from spark_executor.core import connection_store, pending_store +from spark_executor.core.connection_store import ConnectionStore +from spark_executor.core.pending_store import PendingStore +from spark_executor.models import Connection +from spark_executor.tools import connections, submit + + +@pytest.fixture(autouse=True) +def _fresh_stores(tmp_path: Path, monkeypatch): + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(connection_store, "store", ConnectionStore()) + monkeypatch.setattr(pending_store, "DEFAULT_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(pending_store, "store", PendingStore()) + connections.store = connection_store.store + submit.conn_store = connection_store.store + submit.pending_store = pending_store.store + submit.job_store = submit.job_store.__class__() # fresh in-memory job store + + +@pytest.fixture +def log_capture(): + """Attach an in-memory sink to loguru so tests can assert on emitted lines.""" + from common.logging import logger + buf = io.StringIO() + handler_id = logger.add(buf, level="DEBUG", format="{level}|{message}") + yield buf + logger.remove(handler_id) + + +def test_save_connection_emits_info_log(log_capture): + connections.save_connection(name="prod", master="yarn") + text = log_capture.getvalue() + assert "INFO" in text + assert "save_connection enter" in text + assert "DEBUG" in text + assert "connection saved" in text + + +def test_prepare_submit_job_emits_debug_and_info(log_capture): + connections.save_connection(name="prod", master="yarn", deploy_mode="cluster") + log_capture.truncate(0); log_capture.seek(0) + submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py", queue="research") + text = log_capture.getvalue() + assert "DEBUG|prepare_submit_job enter" in text + assert "INFO|prepare_submit_job ok" in text + assert "script_path=/tmp/j.py" in text + assert "queue=research" in text + + +def test_get_unknown_pending_job_emits_debug(log_capture): + log_capture.truncate(0); log_capture.seek(0) + import pytest as _pytest + with _pytest.raises(KeyError): + submit.get_pending_job("p_doesnotexist") + text = log_capture.getvalue() + assert "DEBUG|get_pending_job enter" in text + assert "p_doesnotexist" in text From b3eabe1fb6e1d7ae3794c131daefe4e43284efa4 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 24 Jun 2026 15:09:59 +0800 Subject: [PATCH 23/23] feat: logging to file --- common/logging.py | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/common/logging.py b/common/logging.py index 50bc3f2..7122a46 100644 --- a/common/logging.py +++ b/common/logging.py @@ -13,8 +13,10 @@ Levels used in this project: ERROR - raised exceptions (caller will see the traceback) """ import sys - +from pathlib import Path from loguru import logger +Path("data/logs/debug").mkdir(parents=True, exist_ok=True) +Path("data/logs/info").mkdir(parents=True, exist_ok=True) logger.remove() logger.add( @@ -26,4 +28,24 @@ logger.add( "{name}:{function}:{line} - " "{message}" ), + enqueue=True, + colorize=True, ) + +logger.add( + "data/logs/debug/{time:YYYY-MM-DD}.log", + level="DEBUG", + enqueue=True, + retention="30 days", + compression="gz", + colorize=True, +) + +logger.add( + "data/logs/info/{time:YYYY-MM-DD}.log", + level="INFO", + enqueue=True, + retention="30 days", + compression="gz", + encoding="utf-8", +) \ No newline at end of file