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:
@@ -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}")
|
||||
|
||||
Reference in New Issue
Block a user