colabfit-mcp

colabfit-mcp

KLIFF と KLAY を使用して、ColabFit データセットを検索し、MACE 原子間ポテンシャルを学習する MCP サーバーです。

概要

これは、AI アシスタントに次の機能を提供するModel Context Protocol（MCP）サーバーです。

• ColabFit から科学データセットを検索してダウンロードする
• ローカルのハードウェア（GPU または CPU）で MACE 原子間ポテンシャルを学習する
• エネルギー／力の計算を実行し、OpenKIM テストドライバーでモデルを検証する

会話型 AI とローカルの計算環境を橋渡しします。AI エージェントはこのサーバーを通じて、マシン上でデータを検索し、モデルを学習させ、シミュレーションを実行します。

前提条件

• 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 エージェントは、上記と同じコマンドを起動することで、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 の一般ユーザー向けアプリ（Pro アカウント以外を含む）では、開発者向けエージェントランタイムと同じ方法で、ローカルの 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 など）は、上記の 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 ツールは、通常の pip インストールで Docker を使わずに動作します。

ツール

ツール	説明
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 を保存します。1 回の呼び出しで複数の構造を扱え、スーパーセルのサイズ指定に使う任意の 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 アルミニウム構造を緩和し、最終エネルギーとセルパラメーターを報告してください。

KLAY モデルを使ってバルクシリコンのエネルギー計算を実行する Python スニペットを生成してください。

テストドライバーで検証する：

モデルの検証に利用できるテストドライバーを教えてください。

/home/mcpuser/colabfit/models/si_mace/si_mace__MO_000000000000_000 にある私のシリコンモデルで、ElasticConstantsCrystal テストドライバーを実行してください。

銅の FCC モデルで、EquilibriumCrystalStructure テストと VacancyFormationEnergyRelaxationVolumeCrystal テストを実行してください。

ステータスを確認する：

GPU のステータスを確認し、ローカルにあるすべてのモデルとデータセットを一覧表示してください。

エンドツーエンドのワークフロー：

ColabFit で力を含むシリコンのデータセットを検索し、最適なものをダウンロードして MACE モデルをトレーニングし、バルクのダイヤモンド構造シリコンのエネルギーと力を計算してから、ElasticConstantsCrystal と EquilibriumCrystalStructure のテストドライバーを実行してモデルを検証してください。完了したら、弾性定数と平衡格子定数を報告してください。

トレーニングの停止／キャンセル

MCP サーバーは docker compose up ではなく docker compose run で実行されるため、 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 でインストールしたバージョンでは、GPU の検出を Python の detect_device() のみで処理します。PyTorch がホストの GPU を直接認識できるため、シェルラッパーは不要です。

ローカルインストール（Docker なし）

インストール

pip install colabfit-mcp

これにより、search_datasets、check_local_datasets、download_dataset、build_dataset、 および check_status を利用できます。トレーニングと推論には Docker が必要です。依存関係一式 （CUDA、kim-api、PyG ホイール）は 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

データセットは、KLIFF の Dataset.from_huggingface を介して HuggingFace（colabfit/ 組織）から parquet/arrow ファイルとしてダウンロードされ、ローカルにキャッシュされます。モデルは 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 のワーカープロセス数です。CUDA のフォークによるデッドロックを避けるため、0 のままにしてください。
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 サーバーが応答しない：サーバーは HTTP ではなく、stdio トランスポートを使用します。ネットワークポート経由でアクセスするのではなく、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/Å