From 70400d3ed13acd2a7955e92a0f5cd8c7d27ed47e Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 26 Jun 2026 11:21:53 +0800 Subject: [PATCH] feat(result): get_job_result tool for terminal view of Spark job MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `get_job_status` returns YARN state + the raw response blob, so the terminal fields (finalStatus, diagnostics, trackingUrl, startedTime, finishedTime) are buried inside `raw` and not surfaced in a structured form. Add a new tool that parses them. `get_job_result(job_id)` reuses `yarn_client.get_application_status` and extracts: - finalStatus (SUCCEEDED / FAILED / KILLED / UNDEFINED) - diagnostics (YARN final message) - tracking_url (Spark Web UI) - started_time / finished_time (epoch ms) All five fields are optional: running jobs have no `finishedTime`, and older YARN versions (CDH 5 / H2) may omit some fields. Missing fields stay None — never raise. Coexistence with `get_job_status` is intentional: the latter is for polling the running YARN state, the former is the terminal view. Tests: 4 cases — happy path, running job (no finishedTime), bare-minimum raw (all optionals None), unknown job_id raises KeyError. uv run pytest -> 171 passed. --- spark_executor/models.py | 10 +++ spark_executor/server.py | 16 +++++ spark_executor/tools/result.py | 36 ++++++++++ tests/unit/test_result_tool.py | 123 +++++++++++++++++++++++++++++++++ 4 files changed, 185 insertions(+) create mode 100644 spark_executor/tools/result.py create mode 100644 tests/unit/test_result_tool.py diff --git a/spark_executor/models.py b/spark_executor/models.py index ce27faf..b643593 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -24,6 +24,16 @@ class JobStatus(BaseModel): raw: str = Field(default="") +class JobResult(BaseModel): + application_id: str + state: str + final_status: str | None = None + diagnostics: str | None = None + tracking_url: str | None = None + started_time: int | None = None + finished_time: int | None = None + + class SubmitResult(BaseModel): job_id: str application_id: str diff --git a/spark_executor/server.py b/spark_executor/server.py index 4e94e7e..8fe42be 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -37,6 +37,8 @@ from spark_executor.tools.submit import ( prepare_submit_job, ) +from spark_executor.tools.result import get_job_result + app = FastAPI(title="Spark Executor MCP", version="0.0.1", description="Spark Executor MCP Server") @@ -149,6 +151,20 @@ def _get_job_status(req: JobIdRequest): return get_job_status(req.job_id) +@app.post( + "/get_job_result", + summary="Query YARN for a job's terminal result view", + description=( + "Return a terminal-oriented view of a Spark job: final_status, " + "diagnostics, tracking_url, started_time, and finished_time. " + "This is distinct from get_job_status, which is for polling the " + "running YARN state and returns the raw YARN response." + ), +) +def _get_job_result(req: JobIdRequest): + return get_job_result(req.job_id) + + @app.post( "/get_job_logs", summary="Fetch aggregated container logs for a job", diff --git a/spark_executor/tools/result.py b/spark_executor/tools/result.py new file mode 100644 index 0000000..9c7b505 --- /dev/null +++ b/spark_executor/tools/result.py @@ -0,0 +1,36 @@ +# coding=utf-8 +""" +@Time :2026/6/26 +@Author :tao.chen +""" +import json + +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 JobResult + +store = JobStore() + + +def get_job_result(job_id: str) -> JobResult: + logger.debug(f"get_job_result 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, job.yarn_rm_url) + app = json.loads(raw).get("app", {}) + result = JobResult( + application_id=job.application_id, + state=state, + final_status=app.get("finalStatus"), + diagnostics=app.get("diagnostics"), + tracking_url=app.get("trackingUrl"), + started_time=app.get("startedTime"), + finished_time=app.get("finishedTime"), + ) + logger.info( + f"get_job_result ok job_id={job_id} application_id={job.application_id} " + f"state={state} final_status={result.final_status}" + ) + return result diff --git a/tests/unit/test_result_tool.py b/tests/unit/test_result_tool.py new file mode 100644 index 0000000..83549be --- /dev/null +++ b/tests/unit/test_result_tool.py @@ -0,0 +1,123 @@ +# coding=utf-8 +from datetime import datetime +from unittest.mock import patch + +import pytest + +from spark_executor.core.job_store import JobStore +from spark_executor.models import Job +from spark_executor.tools import result + + +def test_result_returns_parsed_fields(): + result.store = JobStore() + result.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", + yarn_rm_url="http://rm:8088", + ) + ) + raw = { + "app": { + "id": "application_1", + "state": "FINISHED", + "finalStatus": "SUCCEEDED", + "diagnostics": "Application completed successfully", + "trackingUrl": "http://nm:8088/proxy/application_1", + "startedTime": 1700000000000, + "finishedTime": 1700000123000, + } + } + import json + + with patch( + "spark_executor.tools.result.get_application_status", + return_value=("FINISHED", json.dumps(raw)), + ) as m: + out = result.get_job_result("abc") + + assert out.application_id == "application_1" + assert out.state == "FINISHED" + assert out.final_status == "SUCCEEDED" + assert out.diagnostics == "Application completed successfully" + assert out.tracking_url == "http://nm:8088/proxy/application_1" + assert out.started_time == 1700000000000 + assert out.finished_time == 1700000123000 + assert m.call_args.args == ("application_1", "http://rm:8088") + + +def test_result_handles_running_job(): + result.store = JobStore() + result.store.put( + Job( + job_id="running", + application_id="application_2", + script_path="/tmp/j.py", + queue="default", + submit_time=datetime(2026, 6, 24), + connection="prod", + yarn_rm_url="http://rm:8088", + ) + ) + raw = { + "app": { + "id": "application_2", + "state": "RUNNING", + "finalStatus": "UNDEFINED", + "trackingUrl": "http://rm:8088/proxy/application_2", + "startedTime": 1700000000000, + } + } + import json + + with patch( + "spark_executor.tools.result.get_application_status", + return_value=("RUNNING", json.dumps(raw)), + ): + out = result.get_job_result("running") + + assert out.state == "RUNNING" + assert out.final_status == "UNDEFINED" + assert out.finished_time is None + + +def test_result_handles_missing_optional_fields(): + result.store = JobStore() + result.store.put( + Job( + job_id="accepted", + application_id="application_3", + script_path="/tmp/j.py", + queue="default", + submit_time=datetime(2026, 6, 24), + connection="prod", + yarn_rm_url="http://rm:8088", + ) + ) + raw = {"app": {"id": "application_3", "state": "ACCEPTED"}} + import json + + with patch( + "spark_executor.tools.result.get_application_status", + return_value=("ACCEPTED", json.dumps(raw)), + ): + out = result.get_job_result("accepted") + + assert out.state == "ACCEPTED" + assert out.application_id == "application_3" + assert out.final_status is None + assert out.diagnostics is None + assert out.tracking_url is None + assert out.started_time is None + assert out.finished_time is None + + +def test_result_raises_keyerror_for_unknown_job(): + result.store = JobStore() + with pytest.raises(KeyError, match="Unknown job_id"): + result.get_job_result("missing")