colabfit-mcp

colabfit-mcp

一个用于发现 ColabFit 数据集并使用 KLIFF 和 KLAY 训练 MACE 原子间势的 MCP 服务。

概览

这是一个 Model Context Protocol（MCP）服务，让 AI 助手能够：

• 搜索并下载来自 ColabFit 的科学数据集
• 在本地硬件（GPU 或 CPU）上训练 MACE 原子间势
• 运行能量 / 力计算，并使用 OpenKIM 测试驱动程序验证模型

它连接了对话式 AI 与本地计算能力 — AI Agent 通过此服务搜索数据、训练 模型并在你的机器上运行模拟。

前置条件

• Docker 和 Docker Compose v2 — 用于容器化服务
• Git — 用于克隆代码库
• make — 用于快速入门命令（可选；下文提供了手动步骤）
• （可选）NVIDIA GPU + 驱动程序 — 用于 GPU 加速训练
• （可选）nvidia-container-toolkit — Docker 访问 GPU 所必需

本地（非 Docker）安装只需 Python 3.10+。请参阅本地安装（无需 Docker）。

设置

快速入门（推荐）

git clone https://github.com/colabfit/colabfit-mcp.git
cd colabfit-mcp

# One-time setup: creates data directories and .env file
make setup

# Build Docker images with your user ID for proper permissions
make build

然后将 MCP 服务注册到你的客户端（见下方的注册 MCP 服务），并重启客户端。AI 客户端连接时，容器会自动启动。

运行 make help 查看所有可用命令。

手动设置

如果你不想使用 Makefile：

1. 配置环境

cp example.env .env
# Edit .env to customize data directory location if desired

2. 创建数据目录

# Default location
mkdir -p ./colabfit_data/models ./colabfit_data/datasets ./colabfit_data/inference_output ./colabfit_data/test_driver_output

# Or custom location (must match COLABFIT_DATA_ROOT in .env)
# mkdir -p /your/custom/path/{models,datasets,inference_output,test_driver_output}

3. 使用用户 ID 映射进行构建

# This ensures the container user matches your host user and selects the right
# Dockerfile for your platform (CPU-only on macOS, GPU on Linux with NVIDIA)
USER_ID=$(id -u) GROUP_ID=$(id -g) ./start.sh build

注册 MCP 服务

start.sh 会自动检测 NVIDIA GPU 是否可用；如果可用，则启用 GPU 直通，否则回退到 CPU。

Claude Code：

claude mcp add colabfit-mcp -- /path/to/colabfit-mcp/start.sh

将 /path/to/colabfit-mcp 替换为此代码库的绝对路径。 然后重启 Claude Code，使新服务生效。

Claude Desktop：

添加到你的 Claude Desktop 配置中（Settings > Developer > Edit Config）：

{
  "mcpServers": {
    "colabfit-mcp": {
      "command": "/path/to/colabfit-mcp/start.sh",
      "args": ["run", "--rm", "-i", "server"]
    }
  }
}

OpenAI Agent（基于 API，非 ChatGPT 应用）：

支持 MCP 的 OpenAI Agent 可以通过 stdio 连接到此服务，启动命令与上文相同。

将此命令用作 MCP 服务入口点：

/path/to/colabfit-mcp/start.sh

如果你的 Agent 框架需要显式指定 command/args 字段，请使用：

{
  "command": "/path/to/colabfit-mcp/start.sh",
  "args": ["run", "--rm", "-i", "server"]
}

注意：

• 适用于支持注册 MCP 服务的 OpenAI API Agent 运行时。
• ChatGPT 消费者应用（包括非 Pro 账户）无法像开发者 Agent 运行时那样注册本地 stdio MCP 服务。
• 将 /path/to/colabfit-mcp 替换为此代码库的绝对路径。

通用 MCP 客户端设置

此服务使用标准 MCP stdio 传输方式，可与任何兼容 MCP 的客户端配合使用。

入口点（pip install 后或在 Docker 容器中）：

colabfit-mcp          # registered console script
# or
python -m colabfit_mcp

使用 mcp-cli 测试：

pip install mcp-cli
mcp-cli run colabfit-mcp -- colabfit-mcp

任何 stdio MCP 客户端（Gemini、OpenAI Agent、Cursor 等）都可以使用与上文 Claude Desktop 相同的 command / args 模式注册此服务。协议是标准化的 — 所有工具都使用 MCP stdio 传输方式，无需 HTTP 服务或开放端口。

Python SDK 客户端示例：

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

params = StdioServerParameters(
    command="/path/to/colabfit-mcp/start.sh",
    args=["run", "--rm", "-i", "server"],
)

async with stdio_client(params) as (read, write):
    async with ClientSession(read, write) as session:
        await session.initialize()
        tools = await session.list_tools()
        result = await session.call_tool("check_status", {})
        print(result)

使用 pip install mcp 安装客户端库。此服务通过 stdio 使用 JSON-RPC 2.0 — 使用手工构造的 JSON 调用原始 subprocess.Popen 无法正常工作；请使用正规的 MCP 客户端库。

注意：训练和推理需要 Docker（依赖项较重）。search_datasets、check_local_datasets、download_dataset、build_dataset 和 check_status 工具无需 Docker，通过普通的 pip install 即可使用。

工具

工具	描述
search_datasets	按文本、元素、属性和软件搜索 ColabFit 数据库
check_local_datasets	扫描本地数据目录中的已下载数据集，并按元素／属性筛选
download_dataset	通过 KLIFF 从 HuggingFace 下载数据集
train_mace	使用 KLIFF 从头训练 MACE 风格的 KLAY 模型
use_model	使用训练好的 KLAY 模型运行能量／力／弛豫计算，或生成 Python 代码片段
check_status	检查 GPU、软件包、磁盘、现有模型和数据集
list_test_drivers	列出可用的 kimvv 测试驱动程序，可按属性关键词筛选
run_test_driver	对训练好的 KLAY 模型运行 kimvv 测试驱动程序；在带时间戳的子目录中保存 structures.extxyz 和 results.json；单次调用支持多个结构，并可选用 repeat 设置超胞尺寸，或用 async_mode 处理耗时较长的驱动程序
check_test_driver_result	检查异步测试驱动程序作业的状态，并在完成后返回内联结果

可用测试驱动程序（kimvv）

测试驱动程序	描述	属性
EquilibriumCrystalStructure	平衡晶格参数和内聚能	晶格常数、内聚能
ElasticConstantsCrystal	零温下的完整弹性常数张量	弹性常数
CrystalStructureAndEnergyVsPressure	晶体结构和能量随压力的变化	能量与压力的关系
GroundStateCrystalStructure	候选结构中能量最低的晶体结构	基态结构
VacancyFormationEnergyRelaxationVolumeCrystal	空位形成能和弛豫体积	空位形成能、弛豫体积
ClusterEnergyAndForces	在非周期性盒子中对原子簇进行 BFGS 弛豫。适用于分子／非周期性模型。	能量、原子力、弛豫后的位置

典型工作流

1. search_datasets — 查找包含所需元素／属性的数据集
2. download_dataset — 从 HuggingFace 下载数据集（会在本地缓存以供重复使用）
3. train_mace — 使用下载的数据训练 MACE 风格的 KLAY 模型
4. use_model — 运行能量／力／弛豫计算，或生成 Python 代码片段
5. run_test_driver — 使用 OpenKIM 风格的属性测试验证模型

示例提示词

注册 MCP 服务后，可在 Claude Code 或 Claude Desktop 中直接使用以下提示词。

探索可用数据：

在 ColabFit 中搜索包含力数据的硅数据集。哪些数据集最适合训练原子间势？

我本地下载了哪些数据集？其中有没有包含铁和应力数据的？

端到端训练：

找一个铜数据集，下载并用它训练一个 MACE 模型。使用默认设置。

我需要一个磷酸锂势函数。搜索 ColabFit 中包含 Li 和 P 的数据集，选出最合适的一个并开始训练。

运行推理：

使用我的模型 /home/mcpuser/colabfit/models/cu_mace/cu_mace__MO_000000000000_000，计算 FCC 结构块体铜的能量和力。

使用我训练好的模型弛豫 FCC 铝结构，并报告最终能量和晶胞参数。

生成一段 Python 代码，用我的 KLAY 模型计算块体硅的能量。

使用测试驱动程序验证：

有哪些可用的测试驱动程序可以验证我的模型？

在路径为 /home/mcpuser/colabfit/models/si_mace/si_mace__MO_000000000000_000 的硅模型上运行 ElasticConstantsCrystal 测试驱动程序。

在我的铜 FCC 模型上运行 EquilibriumCrystalStructure 和 VacancyFormationEnergyRelaxationVolumeCrystal 测试。

检查状态：

检查我的 GPU 状态，并列出本地的所有模型和数据集。

端到端工作流：

在 ColabFit 中搜索包含力数据的硅数据集，下载最好的一个，训练一个 MACE 模型，计算金刚石立方结构块体硅的能量和力，然后运行 ElasticConstantsCrystal 和 EquilibriumCrystalStructure 测试驱动程序来验证模型。完成后报告弹性常数和平衡晶格参数。

停止 / 取消训练

MCP 服务通过 docker compose run（而非 docker compose up）运行，因此，仅执行 docker compose down 无法停止正在运行的训练容器。 使用以下方法停止服务，包括任何正在进行的训练任务。

使用 Makefile

make stop

不使用 Makefile

# Stop all containers belonging to this project (catches both 'up' and 'run' containers)
docker ps -q --filter "label=com.docker.compose.project=colabfit-mcp" | xargs -r docker stop
docker compose down

如果项目目录名不是 colabfit-mcp，请将筛选值替换为你的目录名（转换为小写）。 你可以使用以下命令检查正在运行的容器上的标签：

docker inspect <container-id> --format '{{ index .Config.Labels "com.docker.compose.project" }}'

训练进度会保存到模型 KIM 子目录中的 training.log 文件 （<model_name>__MO_000000000000_000/training.log）。训练中途停止会丢弃任何 尚未完成的轮次；已完成的轮次及其检查点会保留在磁盘上。

监控训练进度

可通过以下方式查看训练输出：

1. 实时容器日志（推荐）

查看实时训练输出：

# Using Makefile
make logs

# Or directly with docker compose
docker compose logs -f server

按 Ctrl+C 退出（训练会在后台继续）。

2. 持久化日志文件

训练日志文件会写入模型 KIM 子目录：

./colabfit_data/models/<model_name>/<model_name>__MO_000000000000_000/training.log

GPU 支持

start.sh 会自动检测你的 GPU：

• 检测到 NVIDIA GPU：使用 compose.nvidia.yaml 覆盖配置启动，通过 nvidia-container-toolkit 启用 CUDA 直通
• 未检测到 NVIDIA GPU：不使用覆盖配置启动；容器会在运行时自动选择可用的最佳设备（MPS 或 CPU）

通过 pip 安装的版本完全使用 Python 进行 GPU 检测（通过 detect_device()）——无需 shell 包装脚本，因为 PyTorch 可以直接访问主机 GPU。

本地安装（不使用 Docker）

安装

pip install colabfit-mcp

这会启用 search_datasets、check_local_datasets、download_dataset、build_dataset 和 check_status。训练和推理需要 Docker——完整的依赖栈 （CUDA、kim-api、PyG wheels）仅通过 Docker 构建提供支持。

注册到 Claude Code

claude mcp add colabfit-mcp -- colabfit-mcp

注册到 Claude Desktop

将以下内容添加到 Claude Desktop 配置中（Settings > Developer > Edit Config）：

{
  "mcpServers": {
    "colabfit-mcp": {
      "command": "colabfit-mcp"
    }
  }
}

数据目录

默认情况下，数据集和模型存储在 ~/colabfit/ 下。可通过以下方式覆盖：

export COLABFIT_DATA_ROOT=/your/preferred/path

每个工具首次写入数据时，都会自动创建相应的子目录。

要求

• Python 3.10+
• CUDA 12.x + NVIDIA 驱动（用于 GPU 训练；没有 CUDA 时可回退到 CPU）

架构

server container
├── MCP server (FastMCP, stdio)
├── KLIFF (dataset loading, training orchestration)
├── KLAY (MACE-style model construction)
└── Training via KLIFF GNNLightningTrainer

数据集以 parquet/arrow 文件的形式从 HuggingFace（colabfit/ 组织）下载，通过 KLIFF 的 Dataset.from_huggingface 获取并缓存到本地。模型是使用 KLAY 构建的 MACE 风格图模型， 并通过 KLIFF 的 Lightning 训练器进行训练。

容器由 Docker Compose 管理：

• server — MCP 服务 + ML 训练

环境变量

变量	默认值	描述
COLABFIT_DATA_ROOT	./colabfit_data	主机端绑定挂载源目录。容器内的数据根目录始终为 /home/mcpuser/colabfit。
USER_ID	1000	容器用户 ID（应与主机用户一致）
GROUP_ID	1000	容器组 ID（应与主机用户一致）
KLIFF_BATCH_SIZE	4	训练批次大小。若出现 OOM，请调低。
KLIFF_NUM_WORKERS	0	DataLoader 工作进程数。保持为 0，以避免 CUDA fork 死锁。
TRAIN_SIZE	0	训练配置数量（0 = 自动按 90% 划分）
VAL_SIZE	0	验证配置数量（0 = 自动按 10% 划分）
KLIFF_DTYPE	float32	训练精度（默认 float32；更高精度请使用 float64）
COLABFIT_BASE_URL	https://materials.colabfit.org	ColabFit API 基础 URL（用于搜索）
COLABFIT_AUTH_USER	mcp-tool	ColabFit API 身份验证用户名（用于搜索）
COLABFIT_AUTH_PASS	mcp-secret	ColabFit API 身份验证密码（用于搜索）

数据存储：

默认情况下，模型和数据集存储在 ./colabfit_data/ 中（相对于项目根目录），便于随项目移植数据。COLABFIT_DATA_ROOT 仅控制主机端的绑定挂载源目录——无论此设置为何，容器内的数据根目录始终为 /home/mcpuser/colabfit。若要使用在项目克隆之间保持不变的固定主机路径，请在 .env 中设置 COLABFIT_DATA_ROOT：

cp example.env .env
# Edit .env and set: COLABFIT_DATA_ROOT=/home/yourusername/ml_data

Host machine                        Docker container
─────────────                       ────────────────
${COLABFIT_DATA_ROOT}/              /home/mcpuser/colabfit/
├── datasets/          ← bind mount →  ├── datasets/
├── models/            ← bind mount →  ├── models/
├── inference_output/  ← bind mount →  ├── inference_output/
└── test_driver_output/← bind mount →  └── test_driver_output/

用户 ID 映射：

USER_ID 和 GROUP_ID 变量可确保容器用户与主机用户一致，从而避免绑定挂载目录的权限问题。Makefile 会自动检测你的 ID；如有需要，也可以在 .env 中覆盖。

要求

完整列表请参阅先决条件。简而言之：容器化服务需要 Docker + Compose v2，本地安装需要 Python 3.10+。

HPC / 集群用户： HPC 系统通常无法使用 Docker。Apptainer（以前称为 Singularity）可以拉取并转换 Docker 镜像（apptainer pull docker://...），但 Docker Compose 生命周期和 start.sh MCP 注册无法直接迁移到 HPC 环境中。原生 Apptainer/Podman 支持是未来计划实现的目标。

故障排除

安装 torch_scatter 时因“torch not found”而失败：在现有 Python 环境（例如 KDP 容器或系统 Python）中安装时，pip 的构建隔离会导致构建过程无法访问已安装的 torch。请使用 --no-build-isolation：

python -m pip install --no-build-isolation torch-scatter

然后重新安装该软件包，以使用新提供的扩展：

pip install -e ".[full]"

容器未检测到 GPU：确保已安装 nvidia-container-toolkit，并且已重启 Docker 守护进程。使用 docker run --rm --gpus all nvidia/cuda:12.8.0-base-ubuntu22.04 nvidia-smi 进行验证。 如果没有 NVIDIA GPU，请使用 ./start.sh，它会自动回退到 CPU。

MCP 服务无响应：该服务使用 stdio 传输，而非 HTTP。必须通过 docker compose run --rm -i server 启动，不能通过网络端口访问。

---

手动使用：使用训练好的 KLAY 模型进行推理

训练完成后，模型目录（train_mace 返回的 model_path）包含 model.pt 和 kliff_graph.param。可直接将它们与 PyTorch 和 KLIFF 搭配使用。

加载并运行模型

import numpy as np
import torch
from torch_scatter import scatter_add
from kliff.dataset import Configuration
from kliff.transforms.configuration_transforms.graphs.generate_graph import RadialGraph
from ase.build import bulk

atoms = bulk("Si", "diamond", a=5.43)

model_dir = "/home/mcpuser/colabfit/models/colabfit_mace/colabfit_mace__MO_000000000000_000"

# Load model (tries TorchScript first, falls back to torch.load)
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
    model = torch.jit.load(f"{model_dir}/model.pt", map_location=device)
except Exception:
    model = torch.load(f"{model_dir}/model.pt", map_location=device, weights_only=False)
model.eval()
model_dtype = next(model.parameters()).dtype  # match training precision (float32 or float64)

# Build graph — read species/cutoff from kliff_graph.param
transform = RadialGraph(species=["Si"], cutoff=5.0, n_layers=1)
config = Configuration(
    cell=atoms.cell.array,
    species=list(atoms.get_chemical_symbols()),
    coords=atoms.get_positions(),
    PBC=list(atoms.get_pbc()),
    energy=0.0,
    forces=np.zeros((len(atoms), 3)),
)
graph = transform(config)

coords = graph.coords.clone().detach().to(model_dtype).to(device).requires_grad_(True)
energy = model(
    species=graph.species.to(device),
    coords=coords,
    edge_index0=graph.edge_index0.to(device),
    contributions=graph.contributions.to(device),
)
print(f"Energy: {energy.sum().item():.4f} eV")

# Forces via autograd
(grad,) = torch.autograd.grad(energy.sum(), coords)
forces = -scatter_add(grad, graph.images.to(device), dim=0)[:len(atoms)]
print(f"Forces (eV/Å):\n{forces.detach().cpu().numpy()}")

使用 ASE 进行几何优化

use_model 工具中的 _KliffInlineCalculator 会将 KLAY 模型封装为 ASE 计算器。在自定义脚本中，可复现相同的模式：

from ase.optimize import BFGS

# (attach _KliffInlineCalculator from use_model module, or replicate the pattern)
opt = BFGS(atoms, trajectory="relax.traj")
opt.run(fmax=0.01)  # converge forces below 0.01 eV/Å