ml-lab-mcp

An MCP (Model Context Protocol) service deployed on a server, letting large models use this machine as a machine learning / reinforcement learning training resource:

• Code sync — sync_repo(repo_dir, ref?) first fetches then fast-forwards the clone on the server, returning commit/branch/dirty files/ahead-behind counts, ensuring what runs is exactly the code the user pushed to GitHub; only does ff, no reset, and reports errors honestly when diverged. get_repo_state only queries, never modifies.

• Submit experiment — submit_experiment runs any shell command as a background job (bash -lc, so login environments like conda/venv take effect), returning a unique job_id; the uv_project parameter makes the command run in the specified uv project's own environment (uv run --project), so different algorithm projects each use their own environment; job metadata records the workdir's git commit/branch/dirty snapshot, so the code version can be verified later.

• Monitor progress — get_job_status(job_id) returns elapsed_seconds (how long it has run), progress_ratio and eta_seconds (how much time remains, linearly extrapolated from the experiment's self-reported step/timestep/episode/epoch progress, or directly passing through the self-reported eta_seconds); get_job_logs shows the tail of the logs; when multiple experiments run in parallel, they are matched one-to-one by job_id, so there is no confusion.

• Training metrics — read_tensorboard(logdir, tag?) directly parses event files to answer "how is training going": first lists each run's scalar tags, then fetches the specified curve (uniformly downsampled, including latest/min/max), without needing to start a TensorBoard process.

• TensorBoard service — start_tensorboard(logdir, port?, uv_project?) starts the web version for humans to view, returning a URL; stop_tensorboard / list_tensorboards manage it.

• Completion notification — experiments can run for hours; two ways to wait for them to finish: ① wait_for_job(job_id, timeout_seconds) server-side long polling, returns immediately when the job ends, or returns the current status on timeout so you can keep waiting — this uses the normal client→server MCP outbound connection, so the machine running Claude does not need a public IP; ② callback_url — after the job ends, the server POSTs the final metadata (retries 3 times) — note this URL must be reachable from the server, so don't point it at a local machine without a public IP; its real use is pointing to push services like ntfy.sh / Bark / Server酱 to push "training complete" to your phone.

• Fetch results — the result location is decided by the caller (written in the submitted command line); use the generic list_files(path) / read_file(path) to fetch by path; the server does not collect or manage result files.

• Terminate and clean up — cancel_job sends SIGTERM to the entire process group (force=True sends SIGKILL instead); delete_job removes the bookkeeping of finished jobs, delete_path recursively deletes the result/log directory specified by the caller (rejects /, the home directory, and the server's bookkeeping root); job metadata is persisted to disk, so history survives service restarts.

• Public authentication — the HTTP transport enforces a Bearer token (MLLAB_AUTH_TOKEN); requests without or with a wrong token get 401.

Quick Start

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

Without setting MLLAB_AUTH_TOKEN, the HTTP service refuses to start (public deployment enforces authentication).

Related MCP server: colab-autopilot

Client Access

Claude Code:

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

Other MCP clients that support streamable HTTP work the same way: point the URL to http://<server-ip>:8000/mcp and include the Authorization: Bearer <token> header on each request. For local debugging you can use stdio (no auth): uv run mcp dev src/ml_lab_mcp/server.py.

Typical Usage Flow (from the large model's perspective, using DRL training as an example)

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)

Directories and Conventions

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

The job process gets the environment variables JOB_ID, JOB_DIR, PROGRESS_FILE. The experiment script writes JSON to $PROGRESS_FILE by convention, and get_job_status will include this progress and estimate the remaining time from it: it recognizes any pair of (step, total_steps), (timestep, total_timesteps), (episode, total_episodes), (epoch, total_epochs) for linear extrapolation; the script can also directly report eta_seconds. Where result files are written is entirely determined by the command-line arguments; see examples/example_experiment.py.

The callback payload is the content of meta.json (job_id, status, exit_code, etc.), and the delivery result is recorded in the callback_status field, which can be verified via get_job_status. Services like ntfy.sh accept any POST body and work without registration: set callback_url to https://ntfy.sh/<your chosen topic name>, install the ntfy app on your phone and subscribe to the same topic to receive push notifications.

Environment Variables

Security Notes

• Authentication is a static Bearer token (compared in constant time). For public deployment, we recommend adding HTTPS: put an nginx/caddy reverse proxy in front to terminate TLS, since sending the token in plaintext over the public internet is insecure.

• By design, a caller holding the token can execute arbitrary commands and read/delete arbitrary files on the server (as the service process's user). Keep the token safe, and consider running the service under a low-privilege dedicated account.

• start_tensorboard binds to 0.0.0.0 by default, and TensorBoard itself has no authentication — anyone who can reach that port on a public machine can see the training metrics. If that bothers you, restrict the port with a firewall, or don't start TensorBoard and instead use read_tensorboard to have the model relay the info, or use an SSH tunnel.

Extension Directions

• GPU scheduling/queuing: add a queue and concurrency limit before JobManager.submit, and set CUDA_VISIBLE_DEVICES for GPU allocation.

• Large file retrieval: read_file returns at most 200 KB per call (use offset for pagination); for large checkpoints, use rsync/scp or set up a separate file download endpoint.

• Multiple tokens / permission levels: in BearerAuthMiddleware, replace the single token with a token table.

• Callback signing: to prevent forgery, add an HMAC signature to the callback request header for the receiver to verify.

Running Tests

uv run pytest