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

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

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

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

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

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

Tests: 356 passed (up from 242).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-07-08 20:03:40 +08:00
co-authored by Claude Fable 5
parent e285dc0f66
commit 6cf68439a2
14 changed files with 687 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
spark_executor/tools/requests.py
+262 -26
View File
@@ -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.",
)