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