From 6cf68439a2829d699a9f5e1bcd8c406f89c39cfe Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 8 Jul 2026 20:03:40 +0800 Subject: [PATCH] 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 --- README.md | 3 + spark_executor/server.py | 165 +++++++++++++-- spark_executor/tools/external_jobs.py | 65 ++++++ spark_executor/tools/kill.py | 8 +- spark_executor/tools/logs.py | 30 ++- spark_executor/tools/requests.py | 288 +++++++++++++++++++++++--- spark_executor/tools/result.py | 5 +- spark_executor/tools/status.py | 5 +- tests/integration/test_mcp_routes.py | 10 + tests/unit/test_external_jobs.py | 133 ++++++++++++ tests/unit/test_kill_tool.py | 8 + tests/unit/test_logs_tool.py | 8 + tests/unit/test_result_tool.py | 8 + tests/unit/test_status_tool.py | 8 + 14 files changed, 687 insertions(+), 57 deletions(-) create mode 100644 spark_executor/tools/external_jobs.py create mode 100644 tests/unit/test_external_jobs.py diff --git a/README.md b/README.md index 63a71d8..8916969 100644 --- a/README.md +++ b/README.md @@ -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 工具 diff --git a/spark_executor/server.py b/spark_executor/server.py index 285c177..3e156c2 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -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>_'. " + "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): diff --git a/spark_executor/tools/external_jobs.py b/spark_executor/tools/external_jobs.py new file mode 100644 index 0000000..eb60329 --- /dev/null +++ b/spark_executor/tools/external_jobs.py @@ -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 diff --git a/spark_executor/tools/kill.py b/spark_executor/tools/kill.py index b8b2928..5d76aec 100644 --- a/spark_executor/tools/kill.py +++ b/spark_executor/tools/kill.py @@ -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 `) 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}") diff --git a/spark_executor/tools/logs.py b/spark_executor/tools/logs.py index ead5097..8ebc582 100644 --- a/spark_executor/tools/logs.py +++ b/spark_executor/tools/logs.py @@ -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}") diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index fdae819..d9fbee7 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -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 ` 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 ` " + "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>_' (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>_'. 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>_'. 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.", + ) diff --git a/spark_executor/tools/result.py b/spark_executor/tools/result.py index f24dd92..ad9a855 100644 --- a/spark_executor/tools/result.py +++ b/spark_executor/tools/result.py @@ -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}") diff --git a/spark_executor/tools/status.py b/spark_executor/tools/status.py index a143999..edcdefe 100644 --- a/spark_executor/tools/status.py +++ b/spark_executor/tools/status.py @@ -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}") diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index 241e054..3b555b8 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -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 diff --git a/tests/unit/test_external_jobs.py b/tests/unit/test_external_jobs.py new file mode 100644 index 0000000..dc67df5 --- /dev/null +++ b/tests/unit/test_external_jobs.py @@ -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" + ) diff --git a/tests/unit/test_kill_tool.py b/tests/unit/test_kill_tool.py index 2c9bfd4..73151f9 100644 --- a/tests/unit/test_kill_tool.py +++ b/tests/unit/test_kill_tool.py @@ -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( diff --git a/tests/unit/test_logs_tool.py b/tests/unit/test_logs_tool.py index b5f2f78..2e58403 100644 --- a/tests/unit/test_logs_tool.py +++ b/tests/unit/test_logs_tool.py @@ -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( diff --git a/tests/unit/test_result_tool.py b/tests/unit/test_result_tool.py index 5ab6dc8..eb04a11 100644 --- a/tests/unit/test_result_tool.py +++ b/tests/unit/test_result_tool.py @@ -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( diff --git a/tests/unit/test_status_tool.py b/tests/unit/test_status_tool.py index 43d1c23..a6236c6 100644 --- a/tests/unit/test_status_tool.py +++ b/tests/unit/test_status_tool.py @@ -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(