diff --git a/README.md b/README.md index eeaabbb..c644a6d 100644 --- a/README.md +++ b/README.md @@ -121,6 +121,7 @@ MCP 客户端需要先执行 `initialize` 握手，拿到 `mcp-session-id` 后 | `get_external_job_status` | 查询**非本服务提交**的外部 YARN application 状态（按 `application_id` + `connection_name`） | | `get_external_job_result` | 查询外部 YARN application 终态结果视图 | | `get_external_job_logs` | 拉取外部 YARN application 的聚合日志 | +| `list_applications` | 列出 YARN 上所有应用（按 `state` / `queue` / `limit` 过滤），绕过 JobStore | | `fetch_url` | 代理 HTTP GET 到集群内网 URL (host 受 `Connection.url_allowlist` glob allowlist 约束, 空则全拒) | ### Files MCP 工具 diff --git a/spark_executor/core/yarn_client.py b/spark_executor/core/yarn_client.py index 0bb0b56..cd45c5a 100644 --- a/spark_executor/core/yarn_client.py +++ b/spark_executor/core/yarn_client.py @@ -121,13 +121,14 @@ def _base_url(yarn_rm_url: str | None) -> str: def _request(method: str, url: str, *, json_body: dict | None = None, - timeout: float = 30.0, verify: bool | str = True, - auth: httpx.Auth | None = None) -> httpx.Response: + params: dict[str, str] | None = None, timeout: float = 30.0, + verify: bool | str = True, auth: httpx.Auth | None = None) -> httpx.Response: headers = {"Accept": "application/json"} logger.debug(f"YARN {method} {url}" + (f" body={json_body}" if json_body else "")) try: resp = httpx.request( - method, url, json=json_body, headers=headers, timeout=timeout, verify=verify, auth=auth + method, url, json=json_body, params=params, headers=headers, + timeout=timeout, verify=verify, auth=auth ) except httpx.HTTPError as exc: logger.error(f"YARN {method} {url} failed: {exc}") @@ -257,3 +258,54 @@ def kill_application(application_id: str, config: YarnClientConfig) -> None: logger.error(f"YARN PUT {url} -> {resp.status_code}: {resp.text[:500]}") raise YarnError(f"YARN kill returned HTTP {resp.status_code}: {resp.text}") logger.info(f"YARN kill {application_id} -> ok") + + +def list_applications( + config: YarnClientConfig, + *, + state: str | None = None, + queue: str | None = None, + limit: int | None = None, +) -> list[dict]: + """List YARN applications, optionally filtered. + + YARN endpoint: GET /ws/v1/cluster/apps?state=...&queue=...&limit=... + + Filters: + - state: YARN application state. Common values: + "NEW", "NEW_SAVING", "SUBMITTED", "ACCEPTED", "RUNNING", + "FINISHED", "FAILED", "KILLED". + Note: "FINISHED" is the umbrella state covering SUCCEEDED/FAILED/KILLED. + - queue: YARN queue name + - limit: cap on number of returned apps (YARN has no pagination; + callers that need a full enumeration should make multiple + calls with state=... filters or accept the cap) + + Returns a list of YARN app dicts (each with id, name, user, queue, + state, finalStatus, applicationType, startedTime, finishedTime, + trackingUrl, progress, etc). Empty list if no apps match. + + Raises YarnError on transport / 4xx / 5xx. + """ + params: dict[str, str] = {} + if state is not None: + params["state"] = state + if queue is not None: + params["queue"] = queue + if limit is not None: + params["limit"] = str(limit) + + url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps" + resp = _request("GET", url, params=params, + verify=config.verify_for_httpx(), + auth=config.auth_for_httpx()) + if resp.status_code == 404: + # No apps match (or RM doesn't support the endpoint) + return [] + if resp.status_code >= 400: + raise YarnError( + f"YARN list applications failed: {resp.status_code} {resp.text[:200]}" + ) + data = resp.json() + apps_container = data.get("apps") or {} + return apps_container.get("app", []) or [] diff --git a/spark_executor/models.py b/spark_executor/models.py index 1fa327a..b389021 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -48,6 +48,28 @@ class FetchUrlResult(BaseModel): truncated: bool = False +class ApplicationSummary(BaseModel): + """A YARN application summary from /ws/v1/cluster/apps. + + Field names are mapped from the YARN JSON keys to clearer + snake_case names by the tool function. Unused YARN fields + (memorySeconds, vcoreSeconds, preemptedResource*, etc.) are + not exposed — the LLM doesn't need them. + """ + application_id: str + name: str + user: str + queue: str + state: str + final_status: str | None = None + application_type: str | None = None + application_tags: str = "" + started_time: int = 0 + finished_time: int = 0 + tracking_url: str | None = None + progress: float | None = None + + class Connection(BaseModel): name: str # Defaults to "yarn" because that's the literal string spark-submit wants diff --git a/spark_executor/server.py b/spark_executor/server.py index 9a48a19..6453ca6 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -21,6 +21,7 @@ from spark_executor.tools.external_jobs import ( get_external_job_logs, get_external_job_status, get_external_job_result, + list_applications, ) from spark_executor.tools.fetch_url import fetch_url from spark_executor.tools.requests import ( @@ -30,6 +31,7 @@ from spark_executor.tools.requests import ( ExternalJobLogsRequest, ExternalJobStatusRequest, ExternalJobResultRequest, + ListApplicationsRequest, FetchUrlRequest, GetJobLogsRequest, JobIdRequest, @@ -346,6 +348,34 @@ def _get_external_job_result(req: ExternalJobResultRequest): return get_external_job_result(req.application_id, req.connection_name) +@app.post( + "/list_applications", + operation_id="list_applications", + summary="List YARN applications on a cluster, optionally filtered", + description=( + "Query YARN's /ws/v1/cluster/apps endpoint through the named " + "Connection, returning a list of ApplicationSummary records. " + "Bypasses the local JobStore — useful for enumerating apps that " + "were not submitted through this service.\n\n" + "**Filters:** state (YARN state, e.g. 'RUNNING', 'FINISHED', " + "'FAILED'), queue (YARN queue name), limit (default 100, max ~10000). " + "YARN has no offset-based pagination, so for large clusters combine " + "state/queue filters to scope the result. The `FINISHED` state " + "covers SUCCEEDED/FAILED/KILLED.\n\n" + "Returns an empty list if no apps match. The Connection's auth_type " + "/ auth_user / auth_password / ssl_verify / ssl_ca_bundle are reused " + "for the request." + ), +) +def _list_applications(req: ListApplicationsRequest): + return list_applications( + req.connection_name, + state=req.state, + queue=req.queue, + limit=req.limit, + ) + + # --- Connection management tools --- @app.post( diff --git a/spark_executor/tools/external_jobs.py b/spark_executor/tools/external_jobs.py index eb60329..1527317 100644 --- a/spark_executor/tools/external_jobs.py +++ b/spark_executor/tools/external_jobs.py @@ -10,8 +10,13 @@ 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.core.yarn_client import ( + YarnClientConfig, + get_application_logs, + get_application_status, + list_applications as list_applications_yarn, # alias to avoid collision +) +from spark_executor.models import ApplicationSummary, JobResult, JobStatus from spark_executor.tools.connections import store as conn_store @@ -63,3 +68,53 @@ def get_external_job_result(application_id: str, connection_name: str) -> JobRes ) logger.info(f"get_external_job_result ok application_id={application_id} connection_name={connection_name} state={state}") return result + + +def list_applications( + connection_name: str, + state: str | None = None, + queue: str | None = None, + limit: int = 100, +) -> list[ApplicationSummary]: + """List YARN applications on the named cluster, optionally filtered. + + Bypasses the local JobStore (this is for apps not submitted through + this service). The YARN ResourceManager REST endpoint + /ws/v1/cluster/apps is queried through the connection's auth/SSL + config. + + Defaults: limit=100 (YARN has no offset-based pagination, so large + clusters should use state/queue filters to scope the result). + """ + logger.debug( + f"list_applications enter connection_name={connection_name} " + f"state={state} queue={queue} limit={limit}" + ) + conn = conn_store.get(connection_name) + if conn is None: + raise KeyError(f"Connection not found: {connection_name}") + config = YarnClientConfig.from_connection(conn) + raw_apps = list_applications_yarn( + config, state=state, queue=queue, limit=limit + ) + summaries = [ + ApplicationSummary( + application_id=app.get("id", ""), + name=app.get("name", ""), + user=app.get("user", ""), + queue=app.get("queue", ""), + state=app.get("state", ""), + final_status=app.get("finalStatus"), + application_type=app.get("applicationType"), + application_tags=app.get("applicationTags", ""), + started_time=app.get("startedTime", 0), + finished_time=app.get("finishedTime", 0), + tracking_url=app.get("trackingUrl"), + progress=app.get("progress"), + ) + for app in raw_apps + ] + logger.info( + f"list_applications ok connection_name={connection_name} count={len(summaries)}" + ) + return summaries diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index efd740b..27f8284 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -475,3 +475,39 @@ class UpdateConnectionRequest(BaseModel): "Example: ['ccam*'] allows ccam1-ccam99." ), ) + + +class ListApplicationsRequest(BaseModel): + connection_name: str = Field( + ..., + description=( + "Name of a saved Connection (see list_connections) pointing at " + "the YARN cluster to query." + ), + ) + state: str | None = Field( + default=None, + description=( + "Optional YARN application state filter. One of: 'NEW', " + "'NEW_SAVING', 'SUBMITTED', 'ACCEPTED', 'RUNNING', 'FINISHED', " + "'FAILED', 'KILLED'. 'FINISHED' is the umbrella state covering " + "SUCCEEDED/FAILED/KILLED. None = no state filter (returns all " + "states up to `limit`)." + ), + ) + queue: str | None = Field( + default=None, + description=( + "Optional YARN queue name filter (e.g. 'default', 'prod'). " + "None = no queue filter." + ), + ) + limit: int = Field( + default=100, + description=( + "Maximum number of applications to return. YARN has no " + "offset-based pagination, so for large clusters use state/queue " + "filters to scope the result. Max 10000 in practice (YARN's own " + "limit on the limit param)." + ), + ) diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index 6b1e091..df69baf 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -64,12 +64,13 @@ def test_seventeen_tool_routes_registered(): assert "/update_pending_job" in paths -def test_twenty_two_tool_routes_registered(): +def test_twenty_three_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", + "/list_applications", "/fetch_url", "/update_connection", ): diff --git a/tests/unit/test_external_jobs.py b/tests/unit/test_external_jobs.py index dc67df5..36e3f11 100644 --- a/tests/unit/test_external_jobs.py +++ b/tests/unit/test_external_jobs.py @@ -5,8 +5,9 @@ from unittest.mock import patch import pytest from spark_executor.core import connection_store -from spark_executor.models import Connection +from spark_executor.models import ApplicationSummary, Connection from spark_executor.tools import connections, external_jobs +from spark_executor.tools.requests import ListApplicationsRequest def _fresh_stores(): @@ -17,6 +18,13 @@ def _fresh_stores(): external_jobs.conn_store = store +@pytest.fixture +def fresh_stores(tmp_path, monkeypatch): + """Reset connection store singletons to an isolated tmp_path.""" + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + _fresh_stores() + + def test_get_external_job_logs_returns_tailed(): _fresh_stores() external_jobs.conn_store.save( @@ -131,3 +139,154 @@ def test_get_external_job_result_raises_when_connection_missing(): external_jobs.get_external_job_result( application_id="application_1", connection_name="missing" ) + + +def test_list_applications_returns_summaries(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[ + { + "id": "application_1", + "name": "app-one", + "user": "alice", + "queue": "default", + "state": "RUNNING", + "finalStatus": "UNDEFINED", + "applicationType": "SPARK", + "applicationTags": "tag1", + "startedTime": 1000, + "finishedTime": 0, + "trackingUrl": "http://rm:8088/proxy/application_1", + "progress": 75.0, + }, + { + "id": "application_2", + "name": "app-two", + "user": "bob", + "queue": "research", + "state": "FINISHED", + "finalStatus": "SUCCEEDED", + "applicationType": "SPARK", + "applicationTags": "", + "startedTime": 2000, + "finishedTime": 3000, + "trackingUrl": "http://rm:8088/proxy/application_2", + "progress": 100.0, + }, + ], + ) as m: + out = external_jobs.list_applications("prod") + assert len(out) == 2 + assert all(isinstance(item, ApplicationSummary) for item in out) + assert out[0].application_id == "application_1" + assert out[1].application_id == "application_2" + args = m.call_args.args + assert args[0].yarn_rm_url == "http://rm:8088" + + +def test_list_applications_passes_state_filter(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[], + ) as m: + external_jobs.list_applications("prod", state="RUNNING") + assert m.call_args.kwargs == {"state": "RUNNING", "queue": None, "limit": 100} + + +def test_list_applications_passes_queue_filter(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[], + ) as m: + external_jobs.list_applications("prod", queue="research") + assert m.call_args.kwargs == {"state": None, "queue": "research", "limit": 100} + + +def test_list_applications_passes_limit_filter(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[], + ) as m: + external_jobs.list_applications("prod", limit=50) + assert m.call_args.kwargs == {"state": None, "queue": None, "limit": 50} + + +def test_list_applications_default_limit_is_100(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[], + ) as m: + external_jobs.list_applications("prod") + assert m.call_args.kwargs == {"state": None, "queue": None, "limit": 100} + assert ListApplicationsRequest(connection_name="prod").limit == 100 + + +def test_list_applications_returns_empty_list_when_no_apps(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[], + ): + out = external_jobs.list_applications("prod") + assert out == [] + + +def test_list_applications_raises_for_missing_connection(fresh_stores): + with pytest.raises(KeyError, match="Connection not found"): + external_jobs.list_applications("missing") + + +def test_list_applications_maps_yarn_json_to_summary(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + yarn_app = { + "id": "application_42", + "name": "mapped-app", + "user": "carol", + "queue": "prod", + "state": "ACCEPTED", + "finalStatus": "UNDEFINED", + "applicationType": "MAPREDUCE", + "applicationTags": "batch", + "startedTime": 12345, + "finishedTime": 0, + "trackingUrl": "http://rm:8088/proxy/application_42", + "progress": 12.5, + } + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[yarn_app], + ): + out = external_jobs.list_applications("prod") + assert len(out) == 1 + summary = out[0] + assert summary.application_id == "application_42" + assert summary.name == "mapped-app" + assert summary.user == "carol" + assert summary.queue == "prod" + assert summary.state == "ACCEPTED" + assert summary.final_status == "UNDEFINED" + assert summary.application_type == "MAPREDUCE" + assert summary.application_tags == "batch" + assert summary.started_time == 12345 + assert summary.finished_time == 0 + assert summary.tracking_url == "http://rm:8088/proxy/application_42" + assert summary.progress == 12.5