# 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

## 快速开始

### 单机

```bash
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:

```bash
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 也用它做前缀):

```bash
# 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` 改:

```toml
[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:

```bash
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 一一对应,模型的用法几乎不变。唯一多出来的是寻址,只有三条:

| 情形 | 规则 |
|---|---|
| 作业 | job_id 带机器前缀(`lab-a:20260902-120000-ab12cd`),原样传回即可自动路由;`get_job_status` / `get_job_logs` / `cancel_job` / `delete_job` / `wait_for_jobs` 都不用再指定机器 |
| 路径 | `read_file` / `list_files` / `delete_path` / `create_download_link` / `read_tensorboard` / `start_tensorboard` 需要 `machine=<名字>`——同一个路径每台机器上都存在,不指定就无从谈起(只配了一台时可省略) |
| 全局 | `get_machine_status` / `list_jobs` / `list_tensorboards` / `get_repo_state` 不传 `machine` 就是查全部;`sync_repo` 默认 `machine="all"`,一次把所有 clone 钉到同一个 commit |

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

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

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

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

```text
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](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 侧只是多写一层薄封装。

## 环境变量

| 变量 | 默认值 | 说明 |
|---|---|---|
| `MLLAB_AUTH_TOKEN` | (必填) | HTTP 鉴权 Bearer token,不设则拒绝启动 |
| `MLLAB_ROOT` | `~/ml-lab` | 作业簿记根目录 |
| `MLLAB_HOST` | `0.0.0.0` | HTTP 绑定地址 |
| `MLLAB_PORT` | `8000` | HTTP 端口 |
| `MLLAB_TRANSPORT` | worker `streamable-http` / hub `stdio` | 另一个值是对方的默认值 |
| `MLLAB_PUBLIC_HOST` | (自动探测) | 拼进 TensorBoard / 下载 URL 里的主机名/IP |
| `MLLAB_MACHINE_NAME` | 主机名 | 这台机器在回复里的名字,应与 hub 名册里的键一致 |
| `MLLAB_MACHINES` | `~/.ml-lab/machines.toml` | **hub** 的机器名册(TOML 或 JSON) |

## 安全说明

- 鉴权是一层静态 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 签名供接收方验证。

## 运行测试

```bash
uv run pytest
```