# spark-executor-mcp

`spark-executor-mcp` 是一个基于 **FastAPI + fastapi-mcp** 的 Python MCP 服务集合，主要用于让 LLM Agent 以受控方式操作 Spark-on-YARN：保存集群连接、写入或选择 PySpark 脚本、准备提交、二次确认提交、查询状态、读取日志和终止作业。

项目同时内置一个受沙箱限制的 `files_mcp` 服务，用于在指定目录内执行文件读写、查询、移动、复制等操作。

## 核心能力

### Spark Executor MCP

- 管理 Spark/YARN 连接配置。
- 将 LLM 生成的 PySpark 代码写入服务端文件系统。
- 以 **两步提交** 方式运行 Spark 作业：
  1. `prepare_submit_job`：只生成待确认提交记录，不执行 `spark-submit`。
  2. `confirm_submit_job`：用户二次确认后才真正调用 `spark-submit`。
- 查询 Spark/YARN 作业状态、终态结果和聚合日志。
- 通过本地 `job_id` 或 YARN `application_id` 查询同一个作业。
- Kill 正在运行的 YARN application。
- 持久化连接、待提交作业和已提交作业记录。

### Files MCP

- 在 `FILES_MCP_ROOT` 指定的沙箱目录内创建、读取、更新、删除文件。
- 支持目录列举、路径状态查询、glob 搜索、移动和复制。
- 路径会经过沙箱校验，避免访问沙箱外文件。
- 文本内容限制为 UTF-8，单文件读写默认有大小上限。

### 可插拔 MCP 服务

- 根应用启动时会自动发现并挂载内置 MCP 服务。
- 可通过 `MCP_SERVICES` 环境变量选择启用哪些服务。
- 新增 MCP 服务只需要提供一个 FastAPI 子应用和 `SERVICE: McpService` 声明。

## 技术栈

| 类别 | 技术 |
| --- | --- |
| 语言 | Python 3.12+ |
| Web 框架 | FastAPI |
| MCP 暴露 | fastapi-mcp |
| 数据模型 | Pydantic v2 |
| HTTP 客户端 | httpx、httpx-kerberos |
| Spark 集成 | `spark-submit`、YARN REST API、`yarn logs` |
| 日志 | loguru |
| 进程管理 | uvicorn、gunicorn |
| 包管理 | uv |
| 测试 | pytest、pytest-mock |
| 容器 | Docker、Docker Compose |

## 快速开始

### 1. 安装依赖

```bash
uv sync
```

项目依赖在 `pyproject.toml` 和 `uv.lock` 中固定，默认 PyPI 镜像配置为清华源。

### 2. 启动服务

开发环境直接运行：

```bash
uv run main.py
```

默认监听：

```text
0.0.0.0:8000
```

健康检查：

```bash
curl http://localhost:8000/health
```

预期返回：

```json
{"status":"ok"}
```

### 3. MCP 端点

服务启动后会挂载内置 MCP 服务：

| 服务 | MCP endpoint | 说明 |
| --- | --- | --- |
| Spark Executor | `http://localhost:8000/spark-executor-mcp` | Spark-on-YARN 作业管理 |
| Files MCP | `http://localhost:8000/files-mcp` | 沙箱文件系统工具 |

MCP 客户端需要先执行 `initialize` 握手，拿到 `mcp-session-id` 后，再调用 `tools/list` 和 `tools/call`。不要跳过握手。

## MCP 工具概览

### Spark Executor 工具

| 工具 | 作用 |
| --- | --- |
| `save_connection` | 保存或更新一个 Spark/YARN 连接配置 |
| `list_connections` | 列出所有连接配置 |
| `get_connection` | 按名称读取连接配置 |
| `delete_connection` | 删除连接配置 |
| `write_job_file` | 将 LLM 已生成的 PySpark 代码写入服务端文件 |
| `read_job_file` | 读取已存在的 PySpark 脚本内容 |
| `update_job_file` | 覆盖更新已存在的 PySpark 脚本 |
| `prepare_submit_job` | 创建待确认提交记录，不执行 `spark-submit` |
| `update_pending_job` | 在提交前更新待确认作业参数 |
| `list_pending_jobs` | 列出待提交/已提交/已取消/失败的 pending 记录 |
| `get_pending_job` | 查看单个 pending 记录 |
| `cancel_pending_job` | 取消还未提交的 pending 记录 |
| `confirm_submit_job` | 真正执行 `spark-submit` |
| `get_job_status` | 查询 YARN 当前状态 |
| `get_job_result` | 查询终态结果视图 |
| `get_job_logs` | 拉取 YARN 聚合日志 |
| `kill_job` | Kill YARN application |

### Files MCP 工具

| 工具 | 作用 |
| --- | --- |
| `create_file` | 在沙箱内创建文本文件 |
| `read_file` | 读取沙箱内文本文件 |
| `update_file` | 覆盖更新文本文件 |
| `delete_file` | 删除文件或目录 |
| `list_dir` | 列出目录内容 |
| `stat` | 查看路径类型、大小、mtime、权限等信息 |
| `search` | 使用 glob pattern 搜索路径 |
| `move_file` | 移动或重命名文件/目录 |
| `copy_file` | 复制文件/目录 |

## Spark 作业提交流程

Spark 作业提交是故意设计成两步，不要绕过。

### 1. 保存连接

先保存一个连接配置，后续提交作业时通过 `connection` 名称引用它。

示例参数：

```json
{
  "name": "prod-yarn",
  "master": "yarn",
  "deploy_mode": "cluster",
  "yarn_rm_url": "http://yarn-rm.example.com:8088",
  "spark_conf": {
    "spark.yarn.queue": "default"
  }
}
```

调用工具：

```text
save_connection
```

### 2. 写入或选择 PySpark 脚本

如果脚本由 LLM 生成，必须先调用：

```text
write_job_file
```

该工具会把代码写入 `SPARK_EXECUTOR_JOBS_DIR`，并返回真实存在于服务端文件系统中的 `script_path`。

如果脚本已经存在，需要确保传入的是容器或服务进程能访问到的真实路径。只存在于 Agent 上下文里的“虚构路径”会被拒绝。

### 3. 准备提交

调用：

```text
prepare_submit_job
```

这个步骤会：

- 读取并快照指定 `Connection` 的 `master`、`deploy_mode`、`spark_conf`、`yarn_rm_url`。
- 创建 `PendingSubmission` 持久化记录。
- 返回 `pending_id`。
- **不会调用 `spark-submit`。**

注意：如果准备后再修改连接配置，已经创建的 pending job 不会被重新定向到新连接。快照就是快照，别让状态到处乱飞。

### 4. 二次确认提交

用户确认后调用：

```text
confirm_submit_job
```

参数只需要 `pending_id`。该步骤才会真正调用 `spark-submit`，成功后会创建 `Job` 记录，并返回本地 `job_id` 和 YARN `application_id`。

### 5. 查询、读日志或终止

后续工具的 `job_id` 参数同时支持：

- 本地 12 位十六进制 `job_id`。
- YARN `application_id`，例如 `application_17400000001_0001`。

可用工具：

```text
get_job_status
get_job_result
get_job_logs
kill_job
```

## 配置项

### 应用配置

| 环境变量 | 默认值 | 说明 |
| --- | --- | --- |
| `SPARK_EXECUTOR_DATA_DIR` | `./data` | 连接、pending job、job 记录和日志的基础目录 |
| `SPARK_EXECUTOR_JOBS_DIR` | `<SPARK_EXECUTOR_DATA_DIR>/jobs` | `write_job_file` 写入 PySpark 脚本的位置 |
| `SPARK_EXECUTOR_LOG_DIR` | `<SPARK_EXECUTOR_DATA_DIR>/logs` | loguru 文件日志目录 |
| `YARN_RESOURCE_MANAGER_URL` | 空 | 当 Connection 或 Job 快照中没有 `yarn_rm_url` 时的 fallback |
| `SPARK_EXECUTOR_LOG_LEVEL` | `DEBUG` | stderr 和 info 日志级别；debug 文件始终记录 DEBUG |
| `SPARK_EXECUTOR_SPARK_SUBMIT_BIN` | `spark-submit` | 调用 Spark 客户端的命令名或 wrapper 路径 |
| `SPARK_EXECUTOR_SSL_VERIFY_DEFAULT` | `true` | YARN REST 默认是否校验证书 |
| `SPARK_EXECUTOR_SSL_CA_BUNDLE_DEFAULT` | 空 | YARN REST 默认 CA bundle |
| `MCP_SERVICES` | 未设置 | 服务白名单；例如 `spark_executor,files`。未设置表示启用全部内置服务，空字符串表示不启用任何服务 |
| `FILES_MCP_ROOT` | Docker 中为 `/app/data/files` | Files MCP 能访问的唯一沙箱根目录 |

### Gunicorn 配置

| 环境变量 | 默认值 | 说明 |
| --- | --- | --- |
| `GUNICORN_BIND` | `0.0.0.0:8000` | 监听地址 |
| `GUNICORN_WORKERS` | `2` | worker 进程数 |
| `GUNICORN_THREADS` | `1` | 每个 worker 的线程数 |
| `GUNICORN_TIMEOUT` | `120` | 请求超时时间，YARN 日志可能较慢 |
| `GUNICORN_GRACEFUL_TIMEOUT` | `30` | 优雅退出超时 |
| `GUNICORN_KEEPALIVE` | `5` | keepalive 秒数 |
| `GUNICORN_LOGLEVEL` | `info` | gunicorn 日志级别 |

> **重要：生产上建议 `GUNICORN_WORKERS=1`，除非你在前面放了 sticky-session 反向代理。**
>
> MCP session 当前存储在 worker 进程内存里。多 worker 下，`initialize` 可能落到 worker A，下一次 `tools/list` 落到 worker B，然后 B 不认识这个 `mcp-session-id`，就会返回 `Session not found` 或 `Invalid or expired session ID`。这不是玄学，是进程内状态被 round-robin 打散了。

## 数据持久化

默认数据目录是 `./data/`，可通过 `SPARK_EXECUTOR_DATA_DIR` 覆盖。

| 路径 | 说明 |
| --- | --- |
| `data/connections.json` | 保存的 `Connection` 记录 |
| `data/pending_jobs.json` | `PendingSubmission` 记录 |
| `data/jobs.json` | 已提交 `Job` 记录 |
| `data/jobs.json.lock` | `jobs.json` 跨进程写入锁 |
| `data/jobs/` | 默认 PySpark 脚本写入目录 |
| `data/logs/debug/YYYY-MM-DD.log` | DEBUG 日志，按天轮转，保留 30 天，gzip 压缩 |
| `data/logs/info/YYYY-MM-DD.log` | INFO 日志，按天轮转，保留 30 天，gzip 压缩 |
| `data/files/` | Docker 默认 Files MCP 沙箱目录 |

`connections.json`、`pending_jobs.json` 和 `jobs.json` 都使用简单 JSON 文件持久化。当前这是务实方案：足够简单，方便排查，也方便未来迁移到 SQLite。

## Docker 部署

### 1. 准备 Hadoop/YARN 配置

在项目根目录创建或挂载：

```text
./hadoop-conf/
```

至少应包含目标集群匹配的：

```text
core-site.xml
yarn-site.xml
hdfs-site.xml
```

这些文件会以只读方式挂载到容器内：

```text
/etc/hadoop/conf
```

### 2. 准备环境变量

复制模板：

```bash
cp .env.example .env
```

按部署环境修改 `.env`，至少确认：

```text
YARN_RESOURCE_MANAGER_URL=
SPARK_EXECUTOR_DATA_DIR=
SPARK_EXECUTOR_JOBS_DIR=
FILES_MCP_ROOT=
GUNICORN_WORKERS=1
```

`GUNICORN_WORKERS=1` 不是装饰品。除非你真的配置了 sticky-session，否则不要把它调大。

### 3. 启动容器

```bash
docker compose up -d --build
```

查看日志：

```bash
docker compose logs -f mcp-server
```

停止服务：

```bash
docker compose down
```

## 测试

运行完整测试：

```bash
uv run pytest
```

运行单个测试文件：

```bash
uv run pytest tests/unit/test_submit_tool.py -v
```

按名称过滤：

```bash
uv run pytest -v -k cancel_pending
```

测试目录分为：

```text
tests/unit/          # 单元测试
tests/integration/   # FastAPI / MCP 集成测试
```

## 目录结构

```text
main.py                      # 根 FastAPI app；lifespan 中发现并挂载 MCP 服务
common/
  config.py                  # 环境变量配置入口
  factory.py                 # FastApiMCP 初始化封装
  logging.py                 # loguru 配置
  mcp_service.py             # MCP 服务注册、发现、过滤
  sql_guard.py               # PySpark/SQL 安全策略
spark_executor/
  server.py                  # Spark Executor FastAPI 子应用和 MCP tool routes
  service.py                 # Spark Executor MCP 服务声明
  models.py                  # Pydantic 数据模型
  core/                      # spark-submit、YARN client、持久化 store、日志解析
  tools/                     # MCP 工具业务函数和请求模型
files_mcp/
  server.py                  # Files MCP FastAPI 子应用和 MCP tool routes
  service.py                 # Files MCP 服务声明
  core/                      # 文件系统操作和路径沙箱校验
  tools/                     # 文件工具业务函数和请求模型
data/                        # 默认运行时数据目录，gitignored
tests/                       # 单元测试和集成测试
Dockerfile                   # 生产镜像
docker-compose.yml           # 本地/生产容器启动模板
gunicorn.conf.py             # gunicorn 运行配置
pyproject.toml               # Python 项目和依赖声明
uv.lock                      # uv 锁文件
```

## 开发维护说明

### 新增一个 MCP 服务

新增服务的约定很简单，不要搞复杂框架：

1. 创建一个 Python package，例如 `my_service/`。
2. 在 package 内提供 FastAPI app。
3. 新增 `my_service/service.py`，导出 `SERVICE: McpService`。
4. 把 `"my_service.service"` 加到 `common/mcp_service.py` 的 `BUILTIN_SERVICES`。
5. 启动时根应用会自动导入、校验并挂载。

示例：

```python
from common.mcp_service import McpService
from my_service import app

SERVICE = McpService(
    name="my_service",
    app=app,
    mount_path="/my-service-mcp",
)
```

### 新增 MCP 工具路由

所有 MCP 工具路由都应使用 **Pydantic body model**，不要用 query 参数传复杂对象。

正确形态：

```python
@app.post("/tool_name", operation_id="tool_name")
def _tool_name(req: ToolNameRequest):
    return tool_name(**req.model_dump())
```

原因很简单：`fastapi-mcp` 的 `tools/call` 会把参数作为 JSON body 传给 FastAPI。复杂 dict 如果走 query 参数，会变成字符串，最后 422。

### 错误处理约定

工具函数内部保持干净：

- 找不到资源时抛 `KeyError`。
- 非法状态或非法参数抛 `ValueError`。
- 路由层由 FastAPI exception handler 转成 HTTP 状态码。

当前约定：

| 异常 | HTTP 状态 |
| --- | --- |
| `KeyError` | 404 |
| `ValueError` | 400 |
| Pydantic validation error | 422 |
| `FileExistsError`（Files MCP） | 409 |

不要在每个 route body 里手写一堆 `try/except`。那是噪音，不是设计。

### 文件头约定

新增 Python 文件时保持当前项目风格：

```python
# coding=utf-8
"""
@Time   :YYYY/M/D
@Author :tao.chen
"""
```

### 日志约定

使用项目统一 logger：

```python
from common.logging import logger
```

不要在业务模块里自己创建一套 logger。

## 注意事项

1. **两步提交是核心设计，不是可选项。** `prepare_submit_job` 只创建 pending 记录，`confirm_submit_job` 才能执行 `spark-submit`。
2. **不要提交虚构脚本路径。** LLM 生成的 PySpark 必须先用 `write_job_file` 写入服务端文件系统。
3. **pending job 快照连接配置。** prepare 后再修改 Connection，不会影响已经创建的 pending job。
4. **MCP 请求需要 session affinity。** gunicorn 多 worker 会打散内存 session；默认部署请设 `GUNICORN_WORKERS=1`。
5. **`job_id` 参数兼容两种 ID。** 本地 `job_id` 和 YARN `application_id` 都可以传。
6. **Files MCP 只能访问沙箱目录。** 想让它访问某个目录，就把那个目录挂载并设置为 `FILES_MCP_ROOT`，不要绕路径校验。
7. **新工具必须使用 Pydantic body model。** 这是 fastapi-mcp 的调用形态决定的，不是个人偏好。

## 当前状态

- Spark Executor MCP：已实现连接管理、两步提交、作业状态/结果/日志/Kill、PySpark 文件写入/读取/更新。
- Files MCP：已实现沙箱文件 CRUD、搜索、移动、复制。
- 持久化：当前使用 JSON 文件，`jobs.json` 带文件锁；后续可迁移到 SQLite。
- 测试：包含单元测试和集成测试，修改后运行 `uv run pytest`。