feat: add external job tools + improve LLM-facing tool descriptions
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
@@ -117,6 +117,9 @@ MCP 客户端需要先执行 `initialize` 握手，拿到 `mcp-session-id` 后
|
||||
| `get_job_result` | 查询终态结果视图 |
|
||||
| `get_job_logs` | 拉取 YARN 聚合日志 |
|
||||
| `kill_job` | Kill YARN application |
|
||||
| `get_external_job_status` | 查询**非本服务提交**的外部 YARN application 状态（按 `application_id` + `connection_name`） |
|
||||
| `get_external_job_result` | 查询外部 YARN application 终态结果视图 |
|
||||
| `get_external_job_logs` | 拉取外部 YARN application 的聚合日志 |
|
||||
|
||||
### Files MCP 工具
|
||||
|
||||
|
||||
+141
-22
@@ -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):
|
||||
@@ -241,12 +281,67 @@ def _get_job_logs(req: GetJobLogsRequest):
|
||||
"**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. "
|
||||
"**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):
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
# coding=utf-8
|
||||
"""
|
||||
@Time :2026/7/8
|
||||
@Author :tao.chen
|
||||
|
||||
Tools for inspecting YARN applications NOT submitted through this service.
|
||||
These bypass the local JobStore and require the caller to supply both the
|
||||
YARN application_id and the name of a saved Connection.
|
||||
"""
|
||||
import json
|
||||
|
||||
from common.logging import logger
|
||||
from spark_executor.core.yarn_client import YarnClientConfig, get_application_status, get_application_logs
|
||||
from spark_executor.models import JobStatus, JobResult
|
||||
from spark_executor.tools.connections import store as conn_store
|
||||
|
||||
|
||||
def get_external_job_logs(application_id: str, connection_name: str, tail_chars: int = 5000) -> str:
|
||||
"""Fetch aggregated container logs for a YARN application by ID."""
|
||||
logger.debug(f"get_external_job_logs enter application_id={application_id} connection_name={connection_name} tail_chars={tail_chars}")
|
||||
conn = conn_store.get(connection_name)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {connection_name}")
|
||||
config = YarnClientConfig.from_connection(conn)
|
||||
full = get_application_logs(application_id, config)
|
||||
tailed = full[-tail_chars:] if len(full) > tail_chars else full
|
||||
logger.info(
|
||||
f"get_external_job_logs ok application_id={application_id} connection_name={connection_name} "
|
||||
f"full_chars={len(full)} returned_chars={len(tailed)}"
|
||||
)
|
||||
return tailed
|
||||
|
||||
|
||||
def get_external_job_status(application_id: str, connection_name: str) -> JobStatus:
|
||||
"""Query YARN for an external application's current status."""
|
||||
logger.debug(f"get_external_job_status enter application_id={application_id} connection_name={connection_name}")
|
||||
conn = conn_store.get(connection_name)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {connection_name}")
|
||||
config = YarnClientConfig.from_connection(conn)
|
||||
state, raw = get_application_status(application_id, config)
|
||||
logger.info(f"get_external_job_status ok application_id={application_id} connection_name={connection_name} state={state}")
|
||||
return JobStatus(application_id=application_id, state=state, raw=raw)
|
||||
|
||||
|
||||
def get_external_job_result(application_id: str, connection_name: str) -> JobResult:
|
||||
"""Query YARN for an external application's terminal result view."""
|
||||
logger.debug(f"get_external_job_result enter application_id={application_id} connection_name={connection_name}")
|
||||
conn = conn_store.get(connection_name)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {connection_name}")
|
||||
config = YarnClientConfig.from_connection(conn)
|
||||
state, raw = get_application_status(application_id, config)
|
||||
app = json.loads(raw).get("app", {})
|
||||
result = JobResult(
|
||||
application_id=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_external_job_result ok application_id={application_id} connection_name={connection_name} state={state}")
|
||||
return result
|
||||
@@ -21,7 +21,13 @@ def kill_job(job_id: str) -> dict[str, str]:
|
||||
logger.debug(f"kill_job enter job_id={job_id}")
|
||||
job = store.get_either(job_id)
|
||||
if job is None:
|
||||
raise _unknown_job_error(job_id)
|
||||
raise _unknown_job_error(
|
||||
job_id,
|
||||
external_tool_hint=(
|
||||
"the YARN CLI (`yarn application -kill <app_id>`) or the YARN "
|
||||
"UI directly — this service has no external kill tool"
|
||||
),
|
||||
)
|
||||
conn = conn_store.get(job.connection)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {job.connection}")
|
||||
|
||||
@@ -11,14 +11,31 @@ from spark_executor.tools.connections import store as conn_store
|
||||
store = JobStore()
|
||||
|
||||
|
||||
def _unknown_job_error(uid: str) -> KeyError:
|
||||
"""Standard "we tried both IDs and found nothing" message.
|
||||
def _unknown_job_error(uid: str, external_tool_hint: str | None = None) -> Exception:
|
||||
"""Build the right error for a not-found job.
|
||||
|
||||
The agent gets this from confirm_submit_job's response:
|
||||
{"job_id": "a1b2c3d4e5f6", "application_id": "application_...", ...}
|
||||
and routinely confuses which to pass here. Spelling out that BOTH
|
||||
IDs were tried (and what they look like) saves a round trip.
|
||||
and routinely confuses which to pass here. We differentiate two cases:
|
||||
|
||||
- `uid` looks like a YARN application_id (starts with 'application_')
|
||||
AND the caller passed an `external_tool_hint`: the YARN app likely
|
||||
exists, this tool just can't serve it because it was not submitted
|
||||
through this service. Raise ValueError (-> 400 via the FastAPI
|
||||
handler) pointing the agent at the right external tool.
|
||||
- Otherwise: no local record of either form of id. Raise KeyError
|
||||
(-> 404). Spelling out what both IDs look like saves a round trip.
|
||||
"""
|
||||
if uid.startswith("application_") and external_tool_hint:
|
||||
return ValueError(
|
||||
f"job_id={uid!r} looks like a YARN application_id (starts with "
|
||||
f"'application_'), but this tool only works for jobs submitted "
|
||||
f"through this MCP service (no local JobStore record). For YARN "
|
||||
f"applications not submitted here, use {external_tool_hint} "
|
||||
f"instead. (If you actually submitted this job through this "
|
||||
f"service, pass the local job_id — it is a 12-char hex like "
|
||||
f"'a1b2c3d4e5f6'.)"
|
||||
)
|
||||
return KeyError(
|
||||
f"No Job found for id={uid!r} (neither as job_id nor as "
|
||||
f"application_id). Pass the job_id from confirm_submit_job's "
|
||||
@@ -37,7 +54,10 @@ 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_either(job_id)
|
||||
if job is None:
|
||||
raise _unknown_job_error(job_id)
|
||||
raise _unknown_job_error(
|
||||
job_id,
|
||||
external_tool_hint="get_external_job_logs(application_id, connection_name, tail_chars)",
|
||||
)
|
||||
conn = conn_store.get(job.connection)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {job.connection}")
|
||||
|
||||
@@ -17,22 +17,123 @@ class EmptyRequest(BaseModel):
|
||||
|
||||
|
||||
class SaveConnectionRequest(BaseModel):
|
||||
name: str
|
||||
master: str
|
||||
deploy_mode: str = "cluster"
|
||||
yarn_rm_url: str | None = None
|
||||
spark_conf: dict[str, str] | None = None
|
||||
ssl_verify: bool | None = None
|
||||
ssl_ca_bundle: str | None = None
|
||||
auth_type: str = "none"
|
||||
auth_user: str | None = None
|
||||
auth_password: str | None = None
|
||||
auth_principal: str | None = None
|
||||
auth_keytab: str | None = None
|
||||
name: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Unique connection name. Referenced by prepare_submit_job.connection "
|
||||
"and get_external_*.connection_name. Saving with an existing name "
|
||||
"overwrites that record."
|
||||
),
|
||||
)
|
||||
master: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Spark master URL. **Required** at the MCP layer (even though "
|
||||
"the underlying Connection model has a default of 'yarn'). "
|
||||
"Other valid values: 'yarn' (default for YARN), 'spark://host:port' "
|
||||
"(Standalone), 'k8s://...', 'mesos://...', 'local[N]' or 'local[*]'. "
|
||||
"The validator rejects 'yarn-cluster' and bare 'http://...' URLs — "
|
||||
"those are common typos. The literal 'yarn' (not 'yarn-cluster') is "
|
||||
"what spark-submit wants for --master."
|
||||
),
|
||||
)
|
||||
deploy_mode: str = Field(
|
||||
default="cluster",
|
||||
description=(
|
||||
"Spark deploy mode. 'cluster' (default, driver runs in YARN) or "
|
||||
"'client' (driver runs where spark-submit is invoked). Most YARN "
|
||||
"production submissions use 'cluster'."
|
||||
),
|
||||
)
|
||||
yarn_rm_url: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"YARN ResourceManager REST base URL, e.g. 'http://rm-host:8088'. "
|
||||
"**Required** for the 3 get_external_* tools to query YARN "
|
||||
"directly. Optional for prepare_submit_job — spark-submit "
|
||||
"discovers the RM via the cluster config when this is unset."
|
||||
),
|
||||
)
|
||||
spark_conf: dict[str, str] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Dict of Spark conf key→value pairs, passed as --conf flags to "
|
||||
"spark-submit. Example: {'spark.executor.memory': '4g', "
|
||||
"'spark.sql.shuffle.partitions': '200'}. None or empty means "
|
||||
"no extra --conf flags."
|
||||
),
|
||||
)
|
||||
ssl_verify: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Whether to verify the YARN RM TLS certificate. None (default) "
|
||||
"falls back to the global setting; explicit True/False overrides "
|
||||
"the global default for this connection. Set False only for "
|
||||
"self-signed dev clusters."
|
||||
),
|
||||
)
|
||||
ssl_ca_bundle: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Absolute path to a CA bundle file for YARN RM TLS verification. "
|
||||
"Only relevant when the RM uses a private CA. Ignored when "
|
||||
"ssl_verify=False."
|
||||
),
|
||||
)
|
||||
auth_type: str = Field(
|
||||
default="none",
|
||||
description=(
|
||||
"Authentication mode for YARN REST calls. One of: 'none' "
|
||||
"(default, no auth header), 'simple' (pseudo-auth, "
|
||||
"auth_user required), 'basic' (HTTP Basic, auth_user + "
|
||||
"auth_password required), 'kerberos' (SPNEGO via the system "
|
||||
"ticket cache — run kinit beforehand)."
|
||||
),
|
||||
)
|
||||
auth_user: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Username for auth_type='simple' or 'basic'. Ignored when "
|
||||
"auth_type='none' or 'kerberos'."
|
||||
),
|
||||
)
|
||||
auth_password: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Password for auth_type='basic'. Ignored otherwise. Sent on "
|
||||
"every YARN REST request — store with care."
|
||||
),
|
||||
)
|
||||
auth_principal: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Kerberos principal (e.g. 'user@REALM'). Display/audit only; "
|
||||
"the actual SPNEGO handshake uses the system ticket cache. "
|
||||
"Run `kinit <principal>` on the host before invoking the tools."
|
||||
),
|
||||
)
|
||||
auth_keytab: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Absolute path to a Kerberos keytab file. Optional convenience "
|
||||
"for 'kinit -kt' workflows. The service does NOT auto-initialize "
|
||||
"from the keytab — you must `kinit -kt <auth_keytab> <auth_principal>` "
|
||||
"yourself before calling the tools."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class PrepareSubmitJobRequest(BaseModel):
|
||||
connection: str
|
||||
connection: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Name of a saved Connection (call save_connection first, or "
|
||||
"list_connections to see available names). The Connection's "
|
||||
"master / deploy_mode / spark_conf / yarn_rm_url are snapshotted "
|
||||
"into the pending submission at prepare time, so editing the "
|
||||
"Connection afterwards does NOT retarget this pending job."
|
||||
),
|
||||
)
|
||||
app_name: str = Field(
|
||||
...,
|
||||
description="Human-readable application name for tracking the pending submission.",
|
||||
@@ -72,34 +173,115 @@ class PrepareSubmitJobRequest(BaseModel):
|
||||
|
||||
|
||||
class PendingIdRequest(BaseModel):
|
||||
pending_id: str
|
||||
pending_id: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"ID of a pending submission. Format: 'p_' + 12 hex chars, "
|
||||
"e.g. 'p_a1b2c3d4e5f6'. Returned by prepare_submit_job; "
|
||||
"visible via list_pending_jobs."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class UpdatePendingJobRequest(BaseModel):
|
||||
pending_id: str
|
||||
pending_id: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"ID of the pending submission to modify. Format: 'p_' + 12 hex "
|
||||
"chars, e.g. 'p_a1b2c3d4e5f6'. Only PENDING submissions can "
|
||||
"be updated — once SUBMITTED, CANCELLED, or FAILED, the "
|
||||
"pending is terminal."
|
||||
),
|
||||
)
|
||||
script_path: str | None = Field(
|
||||
default=None,
|
||||
description="Optional new absolute path to the PySpark script. If provided, the file must exist and pass SQL guard.",
|
||||
description=(
|
||||
"Optional new absolute path to the PySpark script. If provided, "
|
||||
"the file must exist and pass the SQL guard. Omit to keep the "
|
||||
"current script_path."
|
||||
),
|
||||
)
|
||||
queue: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"New YARN queue name. Omit to keep the current value (PATCH "
|
||||
"semantics: only fields you provide are changed)."
|
||||
),
|
||||
)
|
||||
executor_memory: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"New executor memory, e.g. '4G'. Omit to keep the current value."
|
||||
),
|
||||
)
|
||||
executor_cores: int | None = Field(
|
||||
default=None,
|
||||
description="New cores per executor. Omit to keep the current value.",
|
||||
)
|
||||
num_executors: int | None = Field(
|
||||
default=None,
|
||||
description="New total executor count. Omit to keep the current value.",
|
||||
)
|
||||
app_name: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"New human-readable application name (visible in YARN UI). "
|
||||
"Omit to keep the current value."
|
||||
),
|
||||
)
|
||||
extra_args: dict[str, str] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Replacement dict of additional spark-submit flags (e.g. "
|
||||
"{'jars': '/path/to.jar'}). Unlike the scalar fields, providing "
|
||||
"this REPLACES the entire dict — it is not deep-merged. Omit to "
|
||||
"keep the current value."
|
||||
),
|
||||
)
|
||||
queue: str | None = None
|
||||
executor_memory: str | None = None
|
||||
executor_cores: int | None = None
|
||||
num_executors: int | None = None
|
||||
app_name: str | None = None
|
||||
extra_args: dict[str, str] | None = None
|
||||
|
||||
|
||||
class JobIdRequest(BaseModel):
|
||||
job_id: str
|
||||
job_id: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Either the local job_id (12-char hex, e.g. 'a1b2c3d4e5f6') or "
|
||||
"the YARN application_id (e.g. 'application_17400000001_0001') of "
|
||||
"a job submitted through this service. The lookup tries job_id "
|
||||
"first, then application_id. For YARN applications not submitted "
|
||||
"here, use the get_external_* tools instead."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class GetJobLogsRequest(BaseModel):
|
||||
job_id: str
|
||||
tail_chars: int = 5000
|
||||
job_id: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Either the local job_id (12-char hex, e.g. 'a1b2c3d4e5f6') "
|
||||
"or the YARN application_id (e.g. 'application_17400000001_0001') "
|
||||
"of a job submitted through this service. For YARN applications "
|
||||
"not submitted here, use get_external_job_logs instead."
|
||||
),
|
||||
)
|
||||
tail_chars: int = Field(
|
||||
default=5000,
|
||||
description=(
|
||||
"Return only the last N characters of the aggregated container "
|
||||
"logs. Default 5000. Use a larger value if the head of the log "
|
||||
"(stack traces, driver errors) is being truncated."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ConnectionNameRequest(BaseModel):
|
||||
name: str
|
||||
name: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Name of a saved Connection. Use list_connections to see "
|
||||
"available names. Saving with this name updates an existing "
|
||||
"record (see save_connection)."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class WriteJobFileRequest(BaseModel):
|
||||
@@ -142,3 +324,57 @@ class UpdateJobFileRequest(BaseModel):
|
||||
"Maximum 1 MB to keep the MCP response bounded."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ExternalJobLogsRequest(BaseModel):
|
||||
application_id: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"YARN application_id of the external job, format "
|
||||
"'application_<14-digit-timestamp>_<sequence>' (e.g. "
|
||||
"'application_1740000000001_0001'). This tool is for jobs "
|
||||
"NOT submitted through this MCP service — for those, use "
|
||||
"get_job_logs(job_id=...) instead."
|
||||
),
|
||||
)
|
||||
connection_name: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"Name of a saved Connection (see list_connections) pointing "
|
||||
"at the YARN cluster where the application ran."
|
||||
),
|
||||
)
|
||||
tail_chars: int = Field(
|
||||
default=5000,
|
||||
description="Return only the last N characters of the aggregated container logs.",
|
||||
)
|
||||
|
||||
|
||||
class ExternalJobStatusRequest(BaseModel):
|
||||
application_id: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"YARN application_id of the external job, format "
|
||||
"'application_<14-digit-timestamp>_<sequence>'. Use "
|
||||
"get_job_status(job_id=...) for jobs submitted through this service."
|
||||
),
|
||||
)
|
||||
connection_name: str = Field(
|
||||
...,
|
||||
description="Name of a saved Connection pointing at the YARN cluster.",
|
||||
)
|
||||
|
||||
|
||||
class ExternalJobResultRequest(BaseModel):
|
||||
application_id: str = Field(
|
||||
...,
|
||||
description=(
|
||||
"YARN application_id of the external job, format "
|
||||
"'application_<14-digit-timestamp>_<sequence>'. Use "
|
||||
"get_job_result(job_id=...) for jobs submitted through this service."
|
||||
),
|
||||
)
|
||||
connection_name: str = Field(
|
||||
...,
|
||||
description="Name of a saved Connection pointing at the YARN cluster.",
|
||||
)
|
||||
|
||||
@@ -24,7 +24,10 @@ def get_job_result(job_id: str) -> JobResult:
|
||||
logger.debug(f"get_job_result enter job_id={job_id}")
|
||||
job = store.get_either(job_id)
|
||||
if job is None:
|
||||
raise _unknown_job_error(job_id)
|
||||
raise _unknown_job_error(
|
||||
job_id,
|
||||
external_tool_hint="get_external_job_result(application_id, connection_name)",
|
||||
)
|
||||
conn = conn_store.get(job.connection)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {job.connection}")
|
||||
|
||||
@@ -23,7 +23,10 @@ def get_job_status(job_id: str) -> JobStatus:
|
||||
logger.debug(f"get_job_status enter job_id={job_id}")
|
||||
job = store.get_either(job_id)
|
||||
if job is None:
|
||||
raise _unknown_job_error(job_id)
|
||||
raise _unknown_job_error(
|
||||
job_id,
|
||||
external_tool_hint="get_external_job_status(application_id, connection_name)",
|
||||
)
|
||||
conn = conn_store.get(job.connection)
|
||||
if conn is None:
|
||||
raise KeyError(f"Connection not found: {job.connection}")
|
||||
|
||||
@@ -64,6 +64,16 @@ def test_seventeen_tool_routes_registered():
|
||||
assert "/update_pending_job" in paths
|
||||
|
||||
|
||||
def test_twenty_tool_routes_registered():
|
||||
paths = {r.path for r in app.routes}
|
||||
for path in (
|
||||
"/get_external_job_logs",
|
||||
"/get_external_job_status",
|
||||
"/get_external_job_result",
|
||||
):
|
||||
assert path in paths, f"missing MCP tool route: {path}"
|
||||
|
||||
|
||||
# --- operation_id: pin clean MCP tool names (no auto-generated suffixes) ---
|
||||
#
|
||||
# fastapi-mcp uses each route's OpenAPI `operationId` as the MCP tool name
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
# coding=utf-8
|
||||
import json
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from spark_executor.core import connection_store
|
||||
from spark_executor.models import Connection
|
||||
from spark_executor.tools import connections, external_jobs
|
||||
|
||||
|
||||
def _fresh_stores():
|
||||
"""Reset connection store singletons for a single test."""
|
||||
store = connection_store.ConnectionStore()
|
||||
connection_store.store = store
|
||||
connections.store = store
|
||||
external_jobs.conn_store = store
|
||||
|
||||
|
||||
def test_get_external_job_logs_returns_tailed():
|
||||
_fresh_stores()
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
long_log = "LOG" * 3000
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.get_application_logs",
|
||||
return_value=long_log,
|
||||
) as m:
|
||||
out = external_jobs.get_external_job_logs(
|
||||
application_id="application_1", connection_name="prod", tail_chars=100
|
||||
)
|
||||
assert out == long_log[-100:]
|
||||
args = m.call_args.args
|
||||
assert args[0] == "application_1"
|
||||
assert args[1].yarn_rm_url == "http://rm:8088"
|
||||
|
||||
|
||||
def test_get_external_job_logs_returns_full_when_short():
|
||||
_fresh_stores()
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
short_log = "short log"
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.get_application_logs",
|
||||
return_value=short_log,
|
||||
):
|
||||
out = external_jobs.get_external_job_logs(
|
||||
application_id="application_1", connection_name="prod", tail_chars=5000
|
||||
)
|
||||
assert out == short_log
|
||||
|
||||
|
||||
def test_get_external_job_logs_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
with pytest.raises(KeyError, match="Connection not found"):
|
||||
external_jobs.get_external_job_logs(
|
||||
application_id="application_1", connection_name="missing"
|
||||
)
|
||||
|
||||
|
||||
def test_get_external_job_status_returns_state():
|
||||
_fresh_stores()
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
raw = json.dumps({"app": {"state": "RUNNING"}})
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.get_application_status",
|
||||
return_value=("RUNNING", raw),
|
||||
) as m:
|
||||
out = external_jobs.get_external_job_status(
|
||||
application_id="application_1", connection_name="prod"
|
||||
)
|
||||
assert out.application_id == "application_1"
|
||||
assert out.state == "RUNNING"
|
||||
assert out.raw == raw
|
||||
args = m.call_args.args
|
||||
assert args[0] == "application_1"
|
||||
assert args[1].yarn_rm_url == "http://rm:8088"
|
||||
|
||||
|
||||
def test_get_external_job_status_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
with pytest.raises(KeyError, match="Connection not found"):
|
||||
external_jobs.get_external_job_status(
|
||||
application_id="application_1", connection_name="missing"
|
||||
)
|
||||
|
||||
|
||||
def test_get_external_job_result_parses_app_fields():
|
||||
_fresh_stores()
|
||||
external_jobs.conn_store.save(
|
||||
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
||||
)
|
||||
raw = json.dumps(
|
||||
{
|
||||
"app": {
|
||||
"state": "FINISHED",
|
||||
"finalStatus": "SUCCEEDED",
|
||||
"diagnostics": "",
|
||||
"trackingUrl": "http://rm:8088/proxy/application_1",
|
||||
"startedTime": 100,
|
||||
"finishedTime": 200,
|
||||
}
|
||||
}
|
||||
)
|
||||
with patch(
|
||||
"spark_executor.tools.external_jobs.get_application_status",
|
||||
return_value=("FINISHED", raw),
|
||||
) as m:
|
||||
out = external_jobs.get_external_job_result(
|
||||
application_id="application_1", connection_name="prod"
|
||||
)
|
||||
assert out.application_id == "application_1"
|
||||
assert out.state == "FINISHED"
|
||||
assert out.final_status == "SUCCEEDED"
|
||||
assert out.diagnostics == ""
|
||||
assert out.tracking_url == "http://rm:8088/proxy/application_1"
|
||||
assert out.started_time == 100
|
||||
assert out.finished_time == 200
|
||||
args = m.call_args.args
|
||||
assert args[0] == "application_1"
|
||||
assert args[1].yarn_rm_url == "http://rm:8088"
|
||||
|
||||
|
||||
def test_get_external_job_result_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
with pytest.raises(KeyError, match="Connection not found"):
|
||||
external_jobs.get_external_job_result(
|
||||
application_id="application_1", connection_name="missing"
|
||||
)
|
||||
@@ -63,6 +63,14 @@ def test_kill_job_raises_for_unknown_job():
|
||||
kill.kill_job("missing")
|
||||
|
||||
|
||||
def test_kill_job_raises_400_for_external_application_id(fresh_stores):
|
||||
"""Input that looks like a YARN application_id but is not in the local
|
||||
JobStore must raise ValueError (-> 400) — kill_job has no external
|
||||
equivalent, so the error points to the YARN CLI / UI."""
|
||||
with pytest.raises(ValueError, match="YARN CLI"):
|
||||
kill.kill_job("application_17400000001_0001")
|
||||
|
||||
|
||||
def test_kill_job_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
kill.store.put(
|
||||
|
||||
@@ -68,6 +68,14 @@ def test_get_job_logs_raises_for_unknown_job():
|
||||
logs.get_job_logs("missing")
|
||||
|
||||
|
||||
def test_get_job_logs_raises_400_for_external_application_id(fresh_stores):
|
||||
"""Input that looks like a YARN application_id but is not in the local
|
||||
JobStore must raise ValueError (-> 400) with a hint to use the
|
||||
external tool, NOT a generic KeyError (-> 404)."""
|
||||
with pytest.raises(ValueError, match="get_external_job_logs"):
|
||||
logs.get_job_logs("application_17400000001_0001")
|
||||
|
||||
|
||||
def test_get_job_logs_raises_when_connection_missing(fresh_stores):
|
||||
logs.store.put(
|
||||
Job(
|
||||
|
||||
@@ -136,6 +136,14 @@ def test_result_raises_keyerror_for_unknown_job():
|
||||
assert "application_id" in msg
|
||||
|
||||
|
||||
def test_result_raises_400_for_external_application_id(fresh_stores):
|
||||
"""Input that looks like a YARN application_id but is not in the local
|
||||
JobStore must raise ValueError (-> 400) with a hint to use the
|
||||
external tool, NOT a generic KeyError (-> 404)."""
|
||||
with pytest.raises(ValueError, match="get_external_job_result"):
|
||||
result.get_job_result("application_17400000001_0001")
|
||||
|
||||
|
||||
def test_result_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
result.store.put(
|
||||
|
||||
@@ -66,6 +66,14 @@ def test_get_job_status_raises_for_unknown_job():
|
||||
status.get_job_status("missing")
|
||||
|
||||
|
||||
def test_get_job_status_raises_400_for_external_application_id(fresh_stores):
|
||||
"""Input that looks like a YARN application_id but is not in the local
|
||||
JobStore must raise ValueError (-> 400) with a hint to use the
|
||||
external tool, NOT a generic KeyError (-> 404)."""
|
||||
with pytest.raises(ValueError, match="get_external_job_status"):
|
||||
status.get_job_status("application_17400000001_0001")
|
||||
|
||||
|
||||
def test_get_job_status_raises_when_connection_missing():
|
||||
_fresh_stores()
|
||||
status.store.put(
|
||||
|
||||
Reference in New Issue
Block a user