tao.chen/mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add external job tools + improve LLM-facing tool descriptions

Browse Source 
Add 3 new MCP tools for inspecting YARN applications NOT submitted
through this service: get_external_job_logs, get_external_job_status,
get_external_job_result. Each takes application_id + connection_name
and queries YARN directly, bypassing the local JobStore.

- spark_executor/tools/external_jobs.py: 3 tool functions
- spark_executor/tools/requests.py: 3 new Pydantic body models
  (ExternalJobLogsRequest, ExternalJobStatusRequest,
  ExternalJobResultRequest)
- spark_executor/server.py: 3 new POST routes with explicit operation_id
- tests/unit/test_external_jobs.py: 7 unit tests
- tests/integration/test_mcp_routes.py: assert 20 tool routes
- README.md: list the 3 new tools

To make the LLM pick the right tool and not guess at field values,
also:

- Add Pydantic field descriptions for 22 fields across 8 request models
  (SaveConnectionRequest, UpdatePendingJobRequest, GetJobLogsRequest,
  JobIdRequest, PendingIdRequest, ConnectionNameRequest, plus the new
  ExternalJob*Request models).
- Update 12 route descriptions with cross-references, prerequisite
  context, and 400 behavior notes.
- Refactor _unknown_job_error: an input that looks like a YARN
  application_id (starts with 'application_') now returns HTTP 400
  (ValueError) with a hint message naming the right external tool;
  other not-found cases still return 404 (KeyError). This catches the
  common LLM mistake of passing application_id to the internal
  get_job_* / kill_job tools.
- 4 new unit tests for the 400 behavior.

Tests: 356 passed (up from 242).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-07-08 20:03:40 +08:00
co-authored by Claude Fable 5
parent e285dc0f66
commit 6cf68439a2
14 changed files with 687 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
spark_executor/server.py
+142 -23
View File
@@ -16,10 +16,18 @@ from spark_executor.tools.write_job import write_job_file
from spark_executor.tools.job_file import read_job_file, update_job_file
from spark_executor.tools.kill import kill_job
from spark_executor.tools.logs import get_job_logs
from spark_executor.tools.external_jobs import (
get_external_job_logs,
get_external_job_status,
get_external_job_result,
)
from spark_executor.tools.requests import (
ConnectionNameRequest,
EmptyRequest,
WriteJobFileRequest,
ExternalJobLogsRequest,
ExternalJobStatusRequest,
ExternalJobResultRequest,
GetJobLogsRequest,
JobIdRequest,
PendingIdRequest,
@@ -117,7 +125,9 @@ def _prepare_submit_job(req: PrepareSubmitJobRequest):
"Actually invoke spark-submit for the PendingSubmission identified "
"by pending_id. Requires status=PENDING. On success, transitions the "
"pending entry to SUBMITTED and creates a Job record. On failure, "
"marks the entry FAILED and re-raises."
"marks the entry FAILED and re-raises. A FAILED pending can be "
"re-confirmed — it resets to PENDING for a single fresh attempt — so "
"transient failures (e.g. YARN RM was down) are recoverable."
),
)
def _confirm_submit_job(req: PendingIdRequest):
@@ -128,7 +138,13 @@ def _confirm_submit_job(req: PendingIdRequest):
"/list_pending_jobs",
operation_id="list_pending_jobs",
summary="List all pending submissions",
description="Return every PendingSubmission in any status (PENDING, SUBMITTED, CANCELLED, FAILED).",
description=(
"Return every PendingSubmission in any status (PENDING, SUBMITTED, "
"CANCELLED, FAILED). Call this before prepare_submit_job to check if a "
"submission with the same parameters is already in flight, or after a "
"batch of confirm_submit_job calls to inspect the lifecycle of recent "
"submissions."
),
)
def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
return list_pending_jobs()
@@ -138,7 +154,13 @@ def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
"/get_pending_job",
operation_id="get_pending_job",
summary="Get a single pending submission",
description="Return the PendingSubmission identified by pending_id, including its current status and outcome fields.",
description=(
"Return the PendingSubmission identified by pending_id, including its "
"current status and outcome fields. Use this to inspect a pending "
"submission between prepare_submit_job and confirm_submit_job (e.g. "
"to confirm the snapshotted connection), or to read the error field "
"of a FAILED submission before re-confirming."
),
)
def _get_pending_job(req: PendingIdRequest):
return get_pending_job(req.pending_id)
@@ -186,7 +208,13 @@ def _cancel_pending_job(req: PendingIdRequest):
"confirm_submit_job: the local job_id (12-char hex, e.g. "
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
"'application_17400000001_0001'). The lookup is by job_id first, "
"then by application_id."
"then by application_id. **If you pass a YARN application_id and "
"the app is NOT in the local JobStore, this tool returns HTTP 400** "
"(not 404) with a hint message naming the right external tool. "
"**For YARN applications NOT submitted through this service** "
"(no local JobStore record), use "
"`get_external_job_status(application_id, connection_name)` "
"directly — it bypasses the local registry and queries YARN."
),
)
def _get_job_status(req: JobIdRequest):
@@ -206,7 +234,13 @@ def _get_job_status(req: JobIdRequest):
"confirm_submit_job: the local job_id (12-char hex, e.g. "
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
"'application_17400000001_0001'). The lookup is by job_id first, "
"then by application_id."
"then by application_id. **If you pass a YARN application_id and "
"the app is NOT in the local JobStore, this tool returns HTTP 400** "
"(not 404) with a hint message naming the right external tool. "
"**For YARN applications NOT submitted through this service** "
"(no local JobStore record), use "
"`get_external_job_result(application_id, connection_name)` "
"directly — it bypasses the local registry and queries YARN."
),
)
def _get_job_result(req: JobIdRequest):
@@ -225,7 +259,13 @@ def _get_job_result(req: JobIdRequest):
"confirm_submit_job: the local job_id (12-char hex, e.g. "
"'a1b2c3d4e5f6') and the YARN application_id (e.g. "
"'application_17400000001_0001'). The lookup is by job_id first, "
"then by application_id."
"then by application_id. **If you pass a YARN application_id and "
"the app is NOT in the local JobStore, this tool returns HTTP 400** "
"(not 404) with a hint message naming the right external tool. "
"**For YARN applications NOT submitted through this service** "
"(no local JobStore record), use "
"`get_external_job_logs(application_id, connection_name, tail_chars)` "
"directly — it bypasses the local registry and queries YARN."
),
)
def _get_job_logs(req: GetJobLogsRequest):
@@ -240,13 +280,68 @@ def _get_job_logs(req: GetJobLogsRequest):
"PUT state=KILLED to YARN REST API for the job's application_id. "
"**job_id accepts BOTH identifiers** returned by "
"confirm_submit_job: the local job_id (12-char hex) and the YARN "
"application_id. The lookup is by job_id first, then by application_id."
"application_id. The lookup is by job_id first, then by application_id. "
"**If you pass a YARN application_id and the app is NOT in the "
"local JobStore, this tool returns HTTP 400** (not 404) with a "
"hint pointing to the YARN CLI / UI. **This tool only works for "
"jobs submitted through this service**; there is no external "
"equivalent. For YARN applications you did not submit here, use "
"the YARN CLI / UI directly to kill them."
),
)
def _kill_job(req: JobIdRequest):
return kill_job(req.job_id)
# --- External YARN job tools (bypass JobStore) ---
@app.post(
"/get_external_job_logs",
operation_id="get_external_job_logs",
summary="Query YARN logs for an application not submitted through this service",
description=(
"Fetch aggregated container logs for a YARN application using its "
"application_id and a saved Connection. This bypasses the local JobStore, "
"so it works for jobs submitted outside this MCP service. "
"application_id format is 'application_<14-digit-timestamp>_<sequence>'. "
"For jobs submitted via this service, use get_job_logs(job_id=...) instead."
),
)
def _get_external_job_logs(req: ExternalJobLogsRequest):
return get_external_job_logs(req.application_id, req.connection_name, req.tail_chars)
@app.post(
"/get_external_job_status",
operation_id="get_external_job_status",
summary="Query YARN status for an application not submitted through this service",
description=(
"Return the YARN application state and raw REST response for an "
"application using its application_id and a saved Connection. "
"This bypasses the local JobStore, so it works for jobs submitted "
"outside this MCP service. For jobs submitted via this service, "
"use get_job_status(job_id=...) instead."
),
)
def _get_external_job_status(req: ExternalJobStatusRequest):
return get_external_job_status(req.application_id, req.connection_name)
@app.post(
"/get_external_job_result",
operation_id="get_external_job_result",
summary="Query YARN terminal result for an application not submitted through this service",
description=(
"Return a terminal-oriented view (final_status, diagnostics, tracking_url, "
"started_time, finished_time) for a YARN application using its application_id "
"and a saved Connection. This bypasses the local JobStore. "
"For jobs submitted via this service, use get_job_result(job_id=...) instead."
),
)
def _get_external_job_result(req: ExternalJobResultRequest):
return get_external_job_result(req.application_id, req.connection_name)
# --- Connection management tools ---
@app.post(
@@ -254,9 +349,14 @@ def _kill_job(req: JobIdRequest):
operation_id="save_connection",
summary="Save or update a named Spark connection",
description=(
"Upsert a Connection record (master URL, deploy mode, optional YARN RM URL, "
"spark_conf K/V) keyed by name. Used by prepare_submit_job via the "
"connection parameter."
"Upsert a Connection record (master URL, deploy mode, optional YARN "
"RM URL, spark_conf K/V) keyed by name. Referenced by "
"prepare_submit_job via the connection parameter, and by the 3 "
"get_external_* tools via connection_name. **The Connection's "
"yarn_rm_url is required for get_external_* to work** — spark-submit "
"can discover the RM for submissions, but direct YARN REST queries "
"need an explicit URL. Saving with an existing name overwrites the "
"record in place (no version history)."
),
)
def _save_connection(req: SaveConnectionRequest):
@@ -268,7 +368,12 @@ def _save_connection(req: SaveConnectionRequest):
"/list_connections",
operation_id="list_connections",
summary="List all saved Spark connections",
description="Return every Connection in the registry (model_dump form).",
description=(
"Return every Connection in the registry (model_dump form). Call "
"this before save_connection to see existing names (saving with an "
"existing name overwrites), or after save_connection to verify the "
"record you just stored."
),
)
def _list_connections(_req: EmptyRequest = EmptyRequest()):
return list_connections()
@@ -278,7 +383,12 @@ def _list_connections(_req: EmptyRequest = EmptyRequest()):
"/get_connection",
operation_id="get_connection",
summary="Get a single connection by name",
description="Return the Connection record, or 404 if not found.",
description=(
"Return the Connection record, or 404 if not found. Useful to "
"verify a connection was saved correctly, or to inspect the "
"resolved yarn_rm_url / auth_type before submitting a job or "
"calling one of the get_external_* tools."
),
)
def _get_connection(req: ConnectionNameRequest):
return get_connection(req.name)
@@ -288,7 +398,14 @@ def _get_connection(req: ConnectionNameRequest):
"/delete_connection",
operation_id="delete_connection",
summary="Delete a saved connection",
description="Remove a Connection by name. 404 if not found.",
description=(
"Remove a Connection by name. 404 if not found. Deleting a "
"Connection does NOT affect any pending submission or running job "
"that already references it (the connection details are snapshotted "
"at prepare_submit_job time, and YARN holds the live submission "
"state). New prepare_submit_job calls will fail until you re-save "
"the connection with the same name."
),
)
def _delete_connection(req: ConnectionNameRequest):
return delete_connection(req.name)
@@ -301,16 +418,18 @@ def _delete_connection(req: ConnectionNameRequest):
operation_id="write_job_file",
summary="Write LLM-authored PySpark code to disk",
description=(
"Takes a PySpark code string the LLM has already composed in its "
"context and writes it to a timestamped file under "
"SPARK_EXECUTOR_JOBS_DIR (default ./data/jobs/). Returns the absolute "
"path for use as the script_path argument of prepare_submit_job — the "
"two-step pattern means the LLM writes the file, the user can review "
"it (via read_job_file), and only then is the job submitted.\n\n"
"Note: this tool does NOT generate PySpark code. The calling LLM is "
"expected to have already written the code; this tool only persists "
"it. Code is also run through the SQL safety policy (SELECT/INSERT "
"only) before being written — forbidden statements cause a 400."
"Persist PySpark code you've already written in your context to a "
"timestamped file under SPARK_EXECUTOR_JOBS_DIR (default "
"./data/jobs/). Returns the absolute path to pass as the "
"script_path argument of prepare_submit_job. The two-step pattern "
"(write the file, then prepare) means the user can review the "
"file via read_job_file before anything runs.\n\n"
"Prerequisite: you should have already composed the PySpark code "
"in your own context before calling this tool — it only persists "
"code, it does not generate it. Code is run through the SQL safety "
"policy (SELECT/INSERT only) before being written; forbidden "
"statements cause a 400 with details about which line broke the "
"policy."
),
)
def _write_job_file(req: WriteJobFileRequest):