ml-lab-mcp

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

两种角色,同一份代码:每台机器上跑一个 worker(ml-lab-mcp),多台机器时再跑一个 hub(ml-lab-hub)把它们收在同一个 MCP 后面。客户端始终只注册一个 MCP, 不管背后有几台机器;hub 通过每台 worker 自己的 HTTP 端口和 token 访问它们, 不新增任何通道,也不需要 SSH。只有一台机器时不用 hub,直接注册那台的 worker 即可。

• 多机并行 — 多台机器的意义就是同时跑:submit_experiment 默认 machine="auto",自动挑全集群最空的那张卡并钉上去,返回里的 placement 说明选了谁、为什么;get_machine_status() 不带参数就是整个集群的一张表 (free_gpus 按空闲显存排好序);wait_for_jobs 一次长轮询多个作业、 各机并发,而不是一台等完再等下一台;compare_experiments 能把不同机器上的 实验放在一起排名——各机各自汇总自己的 event 文件,只有摘要过网,不搬数据

• 机器不齐、卡不齐也能用 — 卡是逐张排名的:某台的 GPU 0 满了、GPU 1/2 空着, 照样选得到那台的 GPU 1。刚提交几秒的作业还没吃显存,那张卡在别人眼里仍是"空的", 所以放置会把已被占位但尚未分配显存的卡排到最后(reserved_by 字段), 连着提交两个实验不会叠在同一张卡上。显存需求已知时传 min_free_mb, 没有一张卡放得下就当场拒绝并报出最空的那张,而不是等训练跑起来再 OOM。 连不上的机器、以及(传了 expect_commit 时)clone 不在该 commit 上的机器, 直接不参与竞选,placement.skipped 逐条说明原因;全都不合格才报错, 且错误信息里每台机器一条理由。查状态一类的只读调用用较短的超时(20 秒), 一台关机的机器不会把整集群的视图拖住一分钟

• 代码版本一致性 — 服务端看不见调用方的机器,所以"远程跑的就是我本机 正在改的代码"是一条必须走完的协议,而不是服务端能独自保证的事: ① 在本机读 git rev-parse HEAD 与 git status --porcelain,脏或未推送就停下; ② sync_repo(repo_dir, require_commit=<本机HEAD>) 把 clone 钉到该 commit 并逐字校验结果(clone 有未提交改动同样判定失败,因为那时它的文件已不是该 commit);服务器没有这个 commit 会明确提示"多半没 push"; ③ submit_experiment(expect_commit=<同一 HEAD>)——对不上就拒绝启动。 这是关键:仅仅事后如实记录是不够的,拿错版本的结果回头改代码只会越改越错。 不传 expect_commit 也能跑,但作业元数据会记下 code_verified=false。 sync_repo 只做 ff、不 reset,分叉时如实报错;get_repo_state 只查不改。 多机时 sync_repo 默认 machine="all":一次把每台机器的 clone 都钉到同一个 commit,all_verified 一眼看清全集群是否在同一份代码上——这也是自动选机器的前提。 首次在一台机器上跑某个项目时服务器还没有 clone,给同一个 sync_repo 传 clone_url=<项目的 git 远端> 即可就地创建(cloned 字段说明是否新建), 不需要单独的安装步骤,新建的 clone 也走同一套校验;若 repo_dir 里已有另一个 项目的 clone,clone_url 对不上会直接拒绝,而不是稀里糊涂地同步

• 提交实验 — submit_experiment 以后台作业方式运行任意 shell 命令(bash -lc, conda/venv 等登录环境生效),返回唯一 job_id;uv_project 参数让命令在指定 uv 项目自己的环境里跑(uv run --project),不同算法项目各用各的环境; 作业元数据里记录 workdir 的 git commit/branch/dirty 快照,事后可查证代码版本; logdir 参数(不传则从命令行的 --logdir/--log-dir/--tensorboard-log 自动识别)把 TensorBoard 日志目录和作业绑定,监控时不必再报路径; gpu 参数经 CUDA_VISIBLE_DEVICES 把作业钉到指定卡上(多机时不传 gpu、 用默认的 machine="auto",hub 会替你挑机器和卡)

• 资源感知 — get_machine_status 报告每张 GPU 的显存/利用率/温度和占用进程 (进程通过继承的 JOB_ID 环境变量归因到具体作业,直接回答"GPU 0 被哪个实验占着"), 以及 CPU 负载、内存、磁盘;提交前先看一眼,挑空闲卡传给 gpu 参数。 只报告不拦截,要不要再塞一个实验由调用方判断

• 监控进度与表现 — get_job_status(job_id) 一次调用同时回答"跑到哪了"和 "跑得怎么样":elapsed_seconds(跑了多久)、progress_ratio 与 eta_seconds (还剩多久,按实验自报的 step/timestep/episode/epoch 进度线性外推,或直接透传 自报的 eta_seconds),外加 metrics——该作业 logdir 里几个关键指标的摘要 (最新值、末段均值、最好成绩、趋势)。日志尾部用 get_job_logs;多个实验并行时 靠 job_id 一一对应,不会混淆

• 训练指标 — read_tensorboard(logdir) 不带 tag 时,把每个 run 的每个 scalar tag都压成一行摘要:latest、mean_last10pct、min/max(带出现的 step)、trend(尾段还在涨/在跌/已走平)、non_finite(NaN/inf 计数)。 无论跑了 200 步还是 20 万步,每个 tag 恒定约 90 字节,一次调用就能看完整个实验—— 比逐个 tag 取原始点列省一个数量级的上下文。带 tag 时才返回曲线,格式是并列的 steps/values 数组(4 位有效数字),可选 smoothing(同 TensorBoard UI 的 EMA 滑块,看噪声大的 RL 回报曲线用)和 since_step(长训练增量轮询,只取新增点)。 全程直接解析 event 文件,无需启动 TensorBoard 进程

• 实验对比 — compare_experiments(job_ids?/logdirs?, tag?) 一次调用回答 "exp7 比 exp6 好吗":所有实验的所有 run 在同一 tag 上打分(末段均值、最新值、 最好成绩及其 step、趋势),按名次排序并给出与榜首的差距;tag 不传则自动挑 各实验共有的最"标题级"指标(reward 优先于 loss),含 loss/error 的 tag 自动 按越小越好排。附 arg_diff——各作业命令行里取值不同的 flag,对训练脚本来说 这通常就是超参差异——以及各自的 commit 和机器。跨机器的实验也能直接比: 各机汇总自己的 event 文件,只有摘要过网

• TensorBoard 服务 — start_tensorboard(logdir, port?, uv_project?) 启动 网页版给人看,返回 URL;stop_tensorboard / list_tensorboards 管理

• 完成通知 — 实验可能运行数小时,两条路等它结束: ① wait_for_job(job_id, timeout_seconds) 服务端长轮询,作业结束立刻返回、 超时则返回当前状态可再续等——走的是客户端→服务器的正常 MCP 出站连接, 运行 Claude 的机器不需要公网;多机时用 hub 的 wait_for_jobs([...]), 一次等多个作业、各机并发,一个超时窗口而不是 N 个; ② callback_url 作业结束后由服务端 POST 最终元数据(重试 3 次)——注意这个 URL 必须从服务器可达,所以别指向没有公网的本机,它的实际用途是指向 ntfy.sh / Bark / Server酱 等推送服务,把"训练完成"推到你手机上

• 列表恒定大小 — 机器越用,历史越长,而模型的上下文不会跟着变长,所以两个 列表都分页而不是全量返回:list_jobs 默认给最新 50 条一行摘要 (id、名字、状态、耗时、卡、commit 和命令行开头 200 字),total 是匹配总数、 truncated 说明还有更多,要某个作业的全部细节用 get_job_status; list_files 默认 1000 条,且遍历本身也有上限(20 万个条目), 避免一个手滑传进来的家目录把服务器堵上几分钟。排序按提交时间戳而非 job_id, 同一秒里提交的一批 sweep 也不会乱序

• 取回结果 — 结果位置由调用方决定(写在提交的命令行里)。小文本用 list_files(path) / read_file(path)(单次 ≤200 KB);checkpoint、视频、 打包结果等大文件用 create_download_link(path):同一主机同一端口签发一条 限时签名 URL(默认 1 小时,最长 24 小时),拿到链接的人用普通 curl/wget/浏览器即可下载,无需 token,支持 HTTP Range 断点续传;签名以 MLLAB_AUTH_TOKEN 为密钥,换 token 即作废所有已签链接。目录先用 submit_experiment 打成 tar 再签链接。服务端不收集、不管理结果文件

• 终止与清理 — 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。

多机

每台机器照上面启动 worker,只是各自给一个名字(hub 用这个名字寻址, 下发的 job_id 也用它做前缀):

# 10.2.3.23
MLLAB_MACHINE_NAME=lab-a MLLAB_AUTH_TOKEN=<token-a> uv run ml-lab-mcp
# 10.2.3.24
MLLAB_MACHINE_NAME=lab-b MLLAB_AUTH_TOKEN=<token-b> uv run ml-lab-mcp

在跑 Claude 的那台机器上(不必是服务器)写一份名册,默认路径 ~/.ml-lab/machines.toml,可用 MLLAB_MACHINES 改:

[lab-a]
url   = "http://10.2.3.23:8000"     # 写成 .../mcp 也行,会自动去掉
token = "<token-a>"

[lab-b]
url   = "http://10.2.3.24:8000"
token = "<token-b>"

名册里是所有机器的 token,chmod 600 ~/.ml-lab/machines.toml。然后注册 hub—— 它默认走 stdio,跑在客户端旁边,自己既不监听端口也不需要 token:

claude mcp add ml-lab -- uv run --project /path/to/ml-lab-mcp ml-lab-hub

要让 hub 也对外服务(比如从多处接入),给它 MLLAB_TRANSPORT=streamable-http 和它自己的 MLLAB_AUTH_TOKEN 即可,接入方式 与 worker 相同。

多机下的寻址规则

hub 的工具与 worker 一一对应,模型的用法几乎不变。唯一多出来的是寻址,只有三条:

机器不必配置相同:GPU 数量、有没有 GPU(没有就按 CPU 负载挑机器)、 各自的 uv 环境(uv_project 逐次指定)、磁盘布局都可以不一样。唯一要求是 同一个项目在各机的路径一致——submit_experiment 的 workdir 是同一个字符串 发给被选中的那台。路径确实不同时,要么显式传 machine=,要么在各机做个软链接 统一路径;传了 expect_commit 的话,没有这个 clone 的机器本来就不会被选中。

create_download_link 和 start_tensorboard 返回的 URL 直指那台机器, hub 不中转字节;URL 里的主机名会被换成名册里那个(worker 自己探测的出口 IP 在 VPN/多网卡下经常是错的,名册里的地址按定义是通的)。

Related MCP server: colab-autopilot

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

下面是多机(hub)的写法;单机时把 machine= 相关的参数去掉即可,其余完全一样。

0. 在本机(不是服务器)先确认代码状态:
     git rev-parse HEAD        → SHA
     git status --porcelain    → 有输出就说明有未提交改动,停下来问用户
   sync_repo(repo_dir="/data/proj", require_commit=SHA)   # 默认 machine="all"
   → all_verified 必须为 true;为 false 就别往下走,machines 里看是哪台不对
   # 某台机器上还没有这个项目时,同一个调用加上 clone_url 即可在那台创建:
   sync_repo(repo_dir="/data/proj", require_commit=SHA,
             clone_url="https://github.com/me/proj.git")
   get_machine_status()   # 全集群:哪台哪张卡空着、别的实验占了多少显存
1. 并行提交几组实验(这正是多机的意义):
   submit_experiment(
       command="python train.py --lr 1e-4 --total-timesteps 1000000 --logdir /data/proj/runs/exp7",
       workdir="/data/proj",            # 是 git 仓库 → 元数据记录 commit
       expect_commit=SHA,               # 对不上就拒绝启动,不会静默跑错版本
       uv_project="/data/proj",         # 用该项目自己的 uv 环境
       name="ppo lr=1e-4")              # machine 默认 "auto":自动挑最空的卡
   submit_experiment(..., "--lr 3e-4 ... --logdir /data/proj/runs/exp8",
                     name="ppo lr=3e-4")
   → 各返回一个带机器前缀的 job_id,如 "lab-a:20260902-120000-ab12cd";
     placement 字段说明自动选了哪台哪张卡(也可以显式传 machine= 和 gpu=)
2. wait_for_jobs([id7, id8], timeout_seconds=60)   # 一次等两个,各机并发;超时再调
   get_job_status(id7)        # 一次看全:跑了多久、还剩多久、关键指标、还在不在涨
   get_job_logs(id7)          # 看训练日志尾部
   read_tensorboard("/data/proj/runs/exp7", machine="lab-a")     # 全部 tag 的摘要
   read_tensorboard("/data/proj/runs/exp7", machine="lab-a",
                    tag="rollout/ep_rew_mean", smoothing=0.9)    # 摘要看不出问题时才取曲线
   start_tensorboard("/data/proj/runs/exp7", machine="lab-a", port=6006)  # 给人一个网页 URL
3. 作业结束(回调通知或轮询到 succeeded/failed)后:
   compare_experiments(job_ids=[id7, id8])   # 谁更好?差在哪个超参?跨机器也能比
   list_files("/data/proj/runs/exp7", machine="lab-a")
   read_file("/data/proj/runs/exp7/metrics.json", machine="lab-a")   # 小文本直接读
   create_download_link("/data/proj/runs/exp7/model.pth", machine="lab-a")
   → 返回限时 URL 和现成的 curl 命令,在本机执行即可从那台机器取回大文件;
     整个目录则先 submit_experiment(machine="lab-a",
       command="tar czf /tmp/exp7.tar.gz -C /data/proj/runs exp7")
4. 不要了就清理(先与用户确认):
   cancel_job(id7, force=True)   # 若还在跑
   delete_job(id7)               # 删簿记
   delete_path("/data/proj/runs/exp7", machine="lab-a")  # 删结果/TensorBoard 日志
   stop_tensorboard(6006, machine="lab-a")

目录与约定

$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            # 其运行日志

meta.json 里的 code_verified 记录该作业启动前是否核对过代码版本 (传了 expect_commit 且核对通过才为 true),事后回看指标时据此判断这组数字 是否可信。logdir 记录该作业的 TensorBoard 日志目录(logdir_source 说明是调用方显式指定的 argument,还是从命令行猜出来的 command), get_job_status 据此附带指标摘要;猜不到时 logdir 为 null, 显式传 logdir 参数即可。

作业进程会拿到环境变量 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 订阅同名主题即可收到推送。

代码结构

src/ml_lab_mcp/
├── server.py            # worker:一台机器的 19 个 MCP 工具 + /files 下载 + /rpc
├── hub.py               # hub:同样 19 个工具,但带寻址,转发给各 worker
├── cluster.py           # hub 的地基:名册、RPC、并发扇出、job_id 前缀、放置决策
├── rpc.py               # worker 侧 /rpc 路由(hub 调用 worker 的方式)
├── job_manager.py       # 作业生命周期、进度/ETA、代码版本校验、文件读写删
├── repo.py              # git clone/sync 与状态
├── resources.py         # GPU/CPU/内存/磁盘快照,进程归因到作业
├── tensorboard_tools.py # event 文件解析、指标摘要、托管 TensorBoard 进程
├── analysis.py          # 跨实验排名与超参 diff(可直接吃现成摘要,故能跨机)
├── transfer.py          # 签名下载链接与 /files 流式端点
└── config.py            # 环境变量与两种角色的默认值

hub 只做寻址与合并,不重复实现任何能力:worker 的工具函数原样挂在 /rpc 上, hub 调它们,再把 job_id 加上机器前缀、把 URL 换成名册里的地址。所以给 worker 加一个工具,hub 侧只是多写一层薄封装。

环境变量

安全说明

• 鉴权是一层静态 Bearer token(常数时间比较)。公网部署建议再加 HTTPS: 在前面放 nginx/caddy 反向代理做 TLS 终结,token 明文过公网是不安全的。

• 按设计,持有 token 的调用方可以在服务器上执行任意命令、读取/删除任意文件 (以服务进程的用户身份)。请妥善保管 token,并考虑用低权限专用账户运行服务。

• create_download_link 签发的 URL 在有效期内不需要 token即可下载 对应的那一个文件——把它发到哪里,哪里就能取这个文件;时效默认 1 小时, 换 MLLAB_AUTH_TOKEN 立即作废全部已签链接。

• hub 不新增任何入站面:它只是这些 worker 的客户端,自身走 stdio 时既不监听端口 也不需要 token。代价是名册文件里集中了所有机器的 token,请 chmod 600; 换掉某台的 MLLAB_AUTH_TOKEN 时记得同步改名册,否则那台会以 401 掉线 (报错会明说是 token 不匹配)。

• start_tensorboard 默认绑定 0.0.0.0,而 TensorBoard 本身没有鉴权—— 公网机器上任何能访问该端口的人都能看到训练指标。介意的话用防火墙限制端口, 或不开 TensorBoard、改用 read_tensorboard 由模型转述,或走 SSH 隧道。

扩展方向

• GPU 排队:gpu 参数已能钉卡、hub 的 machine="auto" 会挑最空的卡,但两者 都只是建议:并发上限与排队仍需调用方自律,要硬性约束就在 JobManager.submit 前加队列(单机)或在 hub 的放置逻辑里加准入判断(全集群)。

• 多 token / 权限分级:BearerAuthMiddleware 里把单 token 换成 token 表即可。

• 异构机器:名册目前只有 url/token,项目路径默认各机相同。若各机路径不一, 可在 Machine 上加一个路径映射,由 cluster.Cluster.call 统一改写。

• 回调签名:如需防伪造,可在回调请求头加 HMAC 签名供接收方验证。

运行测试

uv run pytest