Colabfit MCP

colabfit mcp

概述

这是一个 模型上下文协议（MCP）服务器 这使AI助手能够：

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

它连接了对话式人工智能和本地计算——人工智能代理搜索数据、训练 建模，并通过此服务器在您的计算机上运行模拟。

先决条件

• Docker和Docker Compose v2 --对于容器化服务器
• Git --用于克隆存储库
• 制造 --用于快速启动命令（可选；手动步骤记录如下）
• （可选）NVIDIA GPU+驱动程序 --用于GPU加速训练
• （可选） nvidia容器工具包 --Docker访问GPU所必需的

对于本地（非Docker）安装，只需要Python 3.10+。看 本地安装.

设置

快速入门（推荐）

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存在时启用GPU直通，否则退回CPU。

克劳德代码：

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

替换 /path/to/colabfit-mcp 以及此存储库的绝对路径。 然后重新启动Claude Code以使新服务器生效。

克劳德桌面：

添加到您的Claude桌面配置(Settings > Developer > Edit Config):

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

OpenAI代理（基于API，而非ChatGPT应用程序）：

支持MCP的OpenAI代理可以通过以下方式连接到此服务器 stdio 通过启动上面使用的相同命令。

使用此命令作为MCP服务器入口点：

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

如果您的代理框架需要显式的command/args字段，请使用：

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

笔记：

• 这适用于支持MCP服务器注册的基于OpenAI API的代理运行时。
• ChatGPT消费者应用程序（包括非专业帐户）不提供本地 stdio MCP服务器注册的方式与开发人员代理运行时相同。
• 替换 /path/to/colabfit-mcp 以及此存储库的绝对路径。

通用MCP客户端设置

服务器使用标准MCP stdio 传输并与任何兼容MCP的客户端配合使用。

入口点 （pip安装后或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代理、Cursor等）可以使用相同的方式注册服务器 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--raw上使用JSON-RPC 2.0 subprocess.Popen 手工制作的JSON不起作用；使用适当的MCP客户端库。

注意：Docker是训练和推理所必需的（严重依赖）。这 search_datasets, check_local_datasets, download_dataset, build_dataset，以及 check_status 这些工具可以通过简单的pip安装在没有Docker的情况下工作。

工具

可用测试驱动程序（kimvv）

典型工作流程

1. search_datasets --查找包含所需元素/属性的数据集
2. download_dataset --从HuggingFace下载（本地缓存以供重用）
3. train_mace --在下载的数据上训练MACE风格的KLAY模型
4. use_model --运行energy/forces/relax计算或生成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_00000000000_000中的硅模型上运行ElasticConstantsCrystal测试驱动程序。

在我的铜FCC模型上运行平衡晶体结构和空位形成能量弛豫体积晶体测试。

检查状态：

检查我的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  --format '{{ index .Config.Labels "com.docker.compose.project" }}'

培训进度另存为 training.log 在模型的KIM子目录中 (__MO_000000000000_000/training.log).中途停止训练会丢弃任何 进步时代；已完成的纪元及其检查点保存在磁盘上。

监控培训进度

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

1.实时集装箱日志（推荐）

实时查看培训输出：

# Using Makefile
make logs

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

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

2.持久日志文件

Training在模型的KIM子目录中写入日志文件：

./colabfit_data/models//__MO_000000000000_000/training.log

GPU支持

start.sh 自动检测您的GPU：

• NVIDIA GPU存在：以开头 compose.nvidia.yaml 覆盖，通过nvidia容器工具包实现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 mcp add colabfit-mcp -- colabfit-mcp

在Claude Desktop注册

添加到您的Claude桌面配置(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训练；CPU回退在没有CUDA的情况下工作）

建筑

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

数据集从HuggingFace下载(colabfit/ org）通过KLIFF的拼花地板/箭头文件 Dataset.from_huggingface 并在本地缓存。模型是MACE风格的图形 由KLAY建造，并由KLIFF的闪电训练器训练。

Docker Compose管理的容器：

• 服务器 --MCP服务器+ML培训

环境变量

数据存储：

默认情况下，模型和数据集存储在 ./colabfit_data/ （相对于 项目根），使数据可随项目移植。 COLABFIT_DATA_ROOT 控制 只有 主机侧 绑定挂载源——容器内部数据根始终是 /home/mcpuser/colabfit 而不管该设置如何。使用固定主机位置 跨项目克隆持久化，设置 COLABFIT_DATA_ROOT 在 .env:

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/集群用户： Docker通常在HPC系统上不可用。Apptainer（前身为Singularity）可以提取和转换Docker镜像(apptainer pull docker://...)，但是Docker Compose生命周期和 start.sh MCP注册不会直接转换到HPC环境。本地Apptainer/Podman支持是计划中的未来目标。

故障排除

torch_scatter 安装失败，显示“未找到焊炬”：安装到现有系统中时 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模型运行推理

训练后，模型目录(model_path 返回由 train_mace)包含 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/Å