ml-lab-mcp

一个部署在服务器上的 MCP(Model Context Protocol)服务,让大模型可以把这台机器当作 机器学习/强化学习训练资源:

• 代码同步 — sync_repo(repo_dir, ref?) 先 fetch 再 fast-forward 服务器上的 clone,返回 commit/分支/脏文件/领先落后数,确保跑的就是用户推到 GitHub 的代码; 只做 ff,不 reset,分叉时如实报错。get_repo_state 只查不改
• 提交实验 — submit_experiment 以后台作业方式运行任意 shell 命令(bash -lc, conda/venv 等登录环境生效),返回唯一 job_id;uv_project 参数让命令在指定 uv 项目自己的环境里跑(uv run --project),不同算法项目各用各的环境; 作业元数据里记录 workdir 的 git commit/branch/dirty 快照,事后可查证代码版本
• 监控进度 — get_job_status(job_id) 返回 elapsed_seconds(跑了多久)、 progress_ratio 与 eta_seconds(还剩多久,按实验自报的 step/timestep/episode/epoch 进度线性外推,或直接透传自报的 eta_seconds); get_job_logs 看日志尾部;多个实验并行时靠 job_id 一一对应,不会混淆
• 训练指标 — read_tensorboard(logdir, tag?) 直接解析 event 文件回答 "训练表现如何":先列出各 run 的 scalar tag,再取指定曲线(均匀降采样, 含 latest/min/max),无需启动 TensorBoard 进程
• TensorBoard 服务 — start_tensorboard(logdir, port?, uv_project?) 启动 网页版给人看,返回 URL;stop_tensorboard / list_tensorboards 管理
• 完成通知 — 实验可能运行数小时,两条路等它结束: ① wait_for_job(job_id, timeout_seconds) 服务端长轮询,作业结束立刻返回、 超时则返回当前状态可再续等——走的是客户端→服务器的正常 MCP 出站连接, 运行 Claude 的机器不需要公网; ② callback_url 作业结束后由服务端 POST 最终元数据(重试 3 次)——注意这个 URL 必须从服务器可达,所以别指向没有公网的本机,它的实际用途是指向 ntfy.sh / Bark / Server酱 等推送服务,把"训练完成"推到你手机上
• 取回结果 — 结果位置由调用方决定(写在提交的命令行里),用通用的 list_files(path) / read_file(path) 按路径取回;服务端不收集、不管理结果文件
• 终止与清理 — cancel_job SIGTERM 整个进程组(force=True 改发 SIGKILL); delete_job 删除已结束作业的簿记,delete_path 递归删除调用方指定的结果/日志 目录(拒绝 /、家目录和服务端簿记根);作业元数据落盘,服务重启后历史仍在
• 公网鉴权 — HTTP 传输强制 Bearer token(MLLAB_AUTH_TOKEN),未带或带错 token 的请求一律 401

快速开始

cd ml-lab-mcp
uv sync

# 生成一个 token
export MLLAB_AUTH_TOKEN=$(python3 -c 'import secrets; print(secrets.token_urlsafe(32))')

# 启动服务(默认 0.0.0.0:8000,streamable HTTP,路径 /mcp)
uv run ml-lab-mcp

不设置 MLLAB_AUTH_TOKEN 时 HTTP 服务会拒绝启动(公网部署强制鉴权)。

客户端接入

Claude Code:

claude mcp add --transport http ml-lab http://<server-ip>:8000/mcp \
  --header "Authorization: Bearer <token>"

其他支持 streamable HTTP 的 MCP 客户端同理:URL 指向 http://<server-ip>:8000/mcp, 每个请求带 Authorization: Bearer <token> 头。 本地调试可走 stdio(无鉴权):uv run mcp dev src/ml_lab_mcp/server.py。

典型使用流程(大模型视角,DRL 训练为例)

0. sync_repo(repo_dir="/data/proj", ref="main")
   → 确认返回的 commit 就是用户刚推送的那个;dirty/分叉会如实报告
1. submit_experiment(
       command="python train.py --total-timesteps 1000000 --logdir /data/proj/runs/exp7",
       workdir="/data/proj",            # 是 git 仓库 → 元数据记录 commit
       uv_project="/data/proj",         # 用该项目自己的 uv 环境
       name="ppo baseline",
       callback_url="https://ntfy.sh/my-train-topic")   # 可选:训练完推送到手机
   → 记下返回的 job_id
2. wait_for_job(job_id, timeout_seconds=60)   # 会话内等结束:超时就再调一次续等
   get_job_status(job_id)     # 跑了多久 elapsed_seconds / 还剩多久 eta_seconds
   get_job_logs(job_id)       # 看训练日志尾部
   read_tensorboard("/data/proj/runs/exp7")                          # 列 scalar tag
   read_tensorboard("/data/proj/runs/exp7", tag="rollout/ep_rew_mean")  # 看回报曲线
   start_tensorboard("/data/proj/runs/exp7", port=6006)              # 给人一个网页 URL
3. 作业结束(回调通知或轮询到 succeeded/failed)后:
   list_files("/data/proj/runs/exp7")
   read_file("/data/proj/runs/exp7/metrics.json")
4. 不要了就清理(先与用户确认):
   cancel_job(job_id, force=True)   # 若还在跑
   delete_job(job_id)               # 删簿记
   delete_path("/data/proj/runs/exp7")  # 删结果/TensorBoard 日志
   stop_tensorboard(6006)

目录与约定

$MLLAB_ROOT (默认 ~/ml-lab)
├── jobs/
│   └── <job_id>/             # 仅作业簿记,不存实验结果
│       ├── meta.json         # 命令、uv 项目、git 快照、状态、pid、时间戳、退出码
│       ├── output.log        # stdout+stderr 合并日志
│       └── progress.json     # 实验自己写入的进度(可选约定)
└── tensorboard/
    ├── <port>.json           # 托管 TensorBoard 的 pid/logdir/url
    └── <port>.log            # 其运行日志

作业进程会拿到环境变量 JOB_ID、JOB_DIR、PROGRESS_FILE。实验脚本按约定往 $PROGRESS_FILE 写 JSON,get_job_status 就会带上这份进度,并据此估算剩余时间: 识别 (step, total_steps)、(timestep, total_timesteps)、(episode, total_episodes)、(epoch, total_epochs) 任一对做线性外推;脚本也可以直接自报 eta_seconds。结果文件写到哪里完全由命令行参数决定, 参见 examples/example_experiment.py。

回调 payload 即 meta.json 内容(job_id、status、exit_code 等), 投递结果记录在 callback_status 字段,可通过 get_job_status 查证。 ntfy.sh 这类服务接受任意 POST body,免注册即可用:callback_url 填 https://ntfy.sh/<自选主题名>,手机装 ntfy App 订阅同名主题即可收到推送。

环境变量

变量	默认值	说明
MLLAB_AUTH_TOKEN	(必填)	HTTP 鉴权 Bearer token,不设则拒绝启动
MLLAB_ROOT	~/ml-lab	作业簿记根目录
MLLAB_HOST	0.0.0.0	HTTP 绑定地址
MLLAB_PORT	8000	HTTP 端口
MLLAB_TRANSPORT	streamable-http	或 stdio(本地调试,无鉴权)
MLLAB_PUBLIC_HOST	(自动探测)	拼进 TensorBoard URL 里的主机名/IP

安全说明

• 鉴权是一层静态 Bearer token(常数时间比较)。公网部署建议再加 HTTPS: 在前面放 nginx/caddy 反向代理做 TLS 终结,token 明文过公网是不安全的。
• 按设计,持有 token 的调用方可以在服务器上执行任意命令、读取/删除任意文件 (以服务进程的用户身份)。请妥善保管 token,并考虑用低权限专用账户运行服务。
• start_tensorboard 默认绑定 0.0.0.0,而 TensorBoard 本身没有鉴权—— 公网机器上任何能访问该端口的人都能看到训练指标。介意的话用防火墙限制端口, 或不开 TensorBoard、改用 read_tensorboard 由模型转述,或走 SSH 隧道。

扩展方向

• GPU 调度/排队:在 JobManager.submit 前加队列与并发上限, 设置 CUDA_VISIBLE_DEVICES 做卡位分配。
• 大文件取回:read_file 单次最多 200 KB(可用 offset 分页); 大 checkpoint 建议 rsync/scp 或另起文件下载端点。
• 多 token / 权限分级:BearerAuthMiddleware 里把单 token 换成 token 表即可。
• 回调签名:如需防伪造,可在回调请求头加 HMAC 签名供接收方验证。

运行测试

uv run pytest