Claude
dd197019f9
Two problems with the prior prepare_submit_job flow:
1. The agent could pass a script_path that only existed in its own
context (LLM-generated code not yet on disk) or a path on the host
filesystem that's invisible inside the container. The new check
surfaces this as a 400 with a clear remediation hint instead of
letting spark-submit fail later with an opaque FileNotFoundError -> 500.
2. The container-isolation issue: any path the agent gives is interpreted
inside the container. The two ways a file can legitimately exist there
are (a) generate_job_file(code=...) just wrote it to
SPARK_EXECUTOR_JOBS_DIR (the default ./data/jobs/ is the only
gitignored dir that survives restarts), or (b) a host dir was
mounted via -v. The error message spells both out so an agent can
self-correct.
Implementation:
- submit.py: new _check_script_path() that raises ValueError (-> 400)
when the path is missing, empty, or a directory. Called in both
prepare_submit_job and confirm_submit_job (defense in depth).
- prepare_submit_job checks the connection FIRST (KeyError -> 404)
before the script (ValueError -> 400), so an agent with both problems
sees the more fundamental 'unknown connection' error first.
- server.py / requests.py: route description and Pydantic field
description spell out the generate_job_file pattern so an LLM
reading the tool schema learns the right next step.
8 new tests; 8 existing tests adjusted to create real files (they used
synthetic /tmp/*.py paths that don't exist).
229 lines
7.5 KiB
Python
229 lines
7.5 KiB
Python
# coding=utf-8
|
|
"""
|
|
@Time :2026/6/24
|
|
@Author :tao.chen
|
|
"""
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import JSONResponse
|
|
|
|
from spark_executor.tools.connections import (
|
|
delete_connection,
|
|
get_connection,
|
|
list_connections,
|
|
save_connection,
|
|
)
|
|
from spark_executor.tools.generate import generate_job_file
|
|
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,
|
|
GenerateJobFileRequest,
|
|
GetJobLogsRequest,
|
|
JobIdRequest,
|
|
PendingIdRequest,
|
|
PrepareSubmitJobRequest,
|
|
SaveConnectionRequest,
|
|
)
|
|
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 MCP", 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"}
|
|
|
|
|
|
# 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",
|
|
summary="Prepare a Spark job submission (no spark-submit yet)",
|
|
description=(
|
|
"Snapshot the named Connection's master / deploy_mode / spark_conf / "
|
|
"yarn_rm_url into a PendingSubmission record and persist it. "
|
|
"Does NOT invoke spark-submit. Returns pending_id for use with "
|
|
"confirm_submit_job (the user-second-confirmation step).\n\n"
|
|
"REQUIRED PATTERN for LLM-generated code: call generate_job_file(code=...) "
|
|
"first, then pass the returned script_path here. Direct submission with a "
|
|
"synthetic path (one that only exists in the agent's context) will be "
|
|
"rejected with HTTP 400 — the script must exist inside the container's "
|
|
"filesystem. For pre-existing files, mount the host directory into the "
|
|
"container and pass the in-container path."
|
|
),
|
|
)
|
|
def _prepare_submit_job(req: PrepareSubmitJobRequest):
|
|
return prepare_submit_job(**req.model_dump())
|
|
|
|
|
|
@app.post(
|
|
"/confirm_submit_job",
|
|
summary="Confirm and submit a previously-prepared job",
|
|
description=(
|
|
"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."
|
|
),
|
|
)
|
|
def _confirm_submit_job(req: PendingIdRequest):
|
|
return confirm_submit_job(pending_id=req.pending_id)
|
|
|
|
|
|
@app.post(
|
|
"/list_pending_jobs",
|
|
summary="List all pending submissions",
|
|
description="Return every PendingSubmission in any status (PENDING, SUBMITTED, CANCELLED, FAILED).",
|
|
)
|
|
def _list_pending_jobs(_req: EmptyRequest = EmptyRequest()):
|
|
return list_pending_jobs()
|
|
|
|
|
|
@app.post(
|
|
"/get_pending_job",
|
|
summary="Get a single pending submission",
|
|
description="Return the PendingSubmission identified by pending_id, including its current status and outcome fields.",
|
|
)
|
|
def _get_pending_job(req: PendingIdRequest):
|
|
return get_pending_job(req.pending_id)
|
|
|
|
|
|
@app.post(
|
|
"/cancel_pending_job",
|
|
summary="Cancel a pending submission",
|
|
description=(
|
|
"Flip a PENDING (or already-CANCELLED) PendingSubmission to CANCELLED. "
|
|
"Refuses to cancel entries that are SUBMITTED or FAILED — those are "
|
|
"terminal and must be killed via kill_job instead."
|
|
),
|
|
)
|
|
def _cancel_pending_job(req: PendingIdRequest):
|
|
return cancel_pending_job(req.pending_id)
|
|
|
|
|
|
# --- Spark job tools ---
|
|
|
|
@app.post(
|
|
"/get_job_status",
|
|
summary="Query YARN for a job's current status",
|
|
description=(
|
|
"Return the YARN application state (RUNNING / SUCCEEDED / FAILED / "
|
|
"KILLED / ACCEPTED / NEW / NEW_SAVING / SUBMITTED / etc.) plus the "
|
|
"raw YARN REST response body."
|
|
),
|
|
)
|
|
def _get_job_status(req: JobIdRequest):
|
|
return get_job_status(req.job_id)
|
|
|
|
|
|
@app.post(
|
|
"/get_job_logs",
|
|
summary="Fetch aggregated container logs for a job",
|
|
description=(
|
|
"Pull aggregated logs from the YARN ResourceManager. Returns the last "
|
|
"tail_chars characters (default 5000). Requires yarn.log-aggregation-enable "
|
|
"to be true on the target cluster."
|
|
),
|
|
)
|
|
def _get_job_logs(req: GetJobLogsRequest):
|
|
return get_job_logs(req.job_id, tail_chars=req.tail_chars)
|
|
|
|
|
|
@app.post(
|
|
"/kill_job",
|
|
summary="Kill a running job",
|
|
description="PUT state=KILLED to YARN REST API for the job's application_id.",
|
|
)
|
|
def _kill_job(req: JobIdRequest):
|
|
return kill_job(req.job_id)
|
|
|
|
|
|
# --- Connection management tools ---
|
|
|
|
@app.post(
|
|
"/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."
|
|
),
|
|
)
|
|
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",
|
|
summary="List all saved Spark connections",
|
|
description="Return every Connection in the registry (model_dump form).",
|
|
)
|
|
def _list_connections(_req: EmptyRequest = EmptyRequest()):
|
|
return list_connections()
|
|
|
|
|
|
@app.post(
|
|
"/get_connection",
|
|
summary="Get a single connection by name",
|
|
description="Return the Connection record, or 404 if not found.",
|
|
)
|
|
def _get_connection(req: ConnectionNameRequest):
|
|
return get_connection(req.name)
|
|
|
|
|
|
@app.post(
|
|
"/delete_connection",
|
|
summary="Delete a saved connection",
|
|
description="Remove a Connection by name. 404 if not found.",
|
|
)
|
|
def _delete_connection(req: ConnectionNameRequest):
|
|
return delete_connection(req.name)
|
|
|
|
|
|
# --- LLM-driven PySpark generation (Stage 2) ---
|
|
|
|
@app.post(
|
|
"/generate_job_file",
|
|
summary="Write LLM-generated PySpark code to disk",
|
|
description=(
|
|
"Takes a PySpark code string 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 can produce code, the user can review "
|
|
"the resulting file, and only then is the job submitted."
|
|
),
|
|
)
|
|
def _generate_job_file(req: GenerateJobFileRequest):
|
|
return generate_job_file(req.code)
|