Colabfit MCP

colabfit-mcp

An MCP server for discovering ColabFit datasets and training MACE interatomic potentials using KLIFF and KLAY.

Overview

This is a Model Context Protocol (MCP) server that gives AI assistants the ability to:

• Search and download scientific datasets from ColabFit
• Train MACE interatomic potentials on your local hardware (GPU or CPU)
• Run energy/forces calculations and validate models with OpenKIM test drivers

It bridges conversational AI and local compute — the AI agent searches for data, trains models, and runs simulations on your machine through this server.

Prerequisites

• Docker and Docker Compose v2 — for the containerized server
• Git — for cloning the repository
• make — for the quick-start commands (optional; manual steps are documented below)
• (Optional) NVIDIA GPU + drivers — for GPU-accelerated training
• (Optional) [nvidia-container-toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) — required for Docker to access the GPU

For local (non-Docker) installation, only Python 3.10+ is required. See Local Installation.

Setup

Quick Start (Recommended)

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

Then register the MCP server with your client (see Register the MCP server below) and restart your client. The container starts automatically when your AI client connects.

Run make help to see all available commands.

Manual Setup

If you prefer not to use the Makefile:

1. Configure environment

cp example.env .env

Edit .env to customize data directory location if desired

2. Create data directories

# 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. Build with user ID mapping

# 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

Register the MCP server

start.sh automatically detects NVIDIA GPU availability and enables GPU passthrough when present, falling back to CPU otherwise.

Claude Code:

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

Replace /path/to/colabfit-mcp with the absolute path to this repository. Then restart Claude Code for the new server to take effect.

Claude Desktop:

Add to your Claude Desktop config (Settings > Developer > Edit Config):

{

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

OpenAI Agent (API-based, not ChatGPT app):

OpenAI agents that support MCP can connect to this server over stdio by launching the same command used above.

Use this command as the MCP server entrypoint:

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

If your agent framework requires explicit command/args fields, use:

{

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

Notes:

• This is for OpenAI API-based agent runtimes that support MCP server registration.
• The ChatGPT consumer app (including non-Pro accounts) does not provide local stdio MCP server registration in the same way as developer agent runtimes.
• Replace /path/to/colabfit-mcp with the absolute path to this repository.

Generic MCP Client Setup

The server uses standard MCP stdio transport and works with any MCP-compatible client.

Entry point (after pip install or in the Docker container):

colabfit-mcp          # registered console script

or

python -m colabfit_mcp

Testing with mcp-cli:

pip install mcp-cli

mcp-cli run colabfit-mcp -- colabfit-mcp

Any stdio MCP client (Gemini, OpenAI agents, Cursor, etc.) can register the server using the same command / args pattern as Claude Desktop above. The protocol is standardized — all tools use MCP stdio transport, no HTTP server or open port is required.

Python SDK client example:

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)

Install the client library with pip install mcp. The server uses JSON-RPC 2.0 over stdio — raw subprocess.Popen with hand-crafted JSON will not work; use a proper MCP client library.

Note: Docker is required for training and inference (heavy dependencies). The search_datasets, check_local_datasets, download_dataset, build_dataset, and check_status tools work without Docker via a plain pip install.

Tools

Available Test Drivers (kimvv)

Typical Workflow

1. search_datasets — find datasets with the elements/properties you need
2. download_dataset — download from HuggingFace (cached locally for reuse)
3. train_mace — train a MACE-style KLAY model on the downloaded data
4. use_model — run energy/forces/relax calculations or generate a Python snippet
5. run_test_driver — validate the model against OpenKIM-style property tests

Sample Prompts

The following prompts work directly in Claude Code or Claude Desktop once the MCP server is registered.

Explore available data:

Search ColabFit for silicon datasets that include forces. Which ones look best for training an interatomic potential?

What datasets do I have downloaded locally? Do any contain iron with stress data?

End-to-end training:

Find a dataset for copper, download it, and train a MACE model on it. Use default settings.

I need a potential for lithium phosphate. Search ColabFit for Li and P datasets, pick the most suitable one, and start training.

Run inference:

Use my model at /home/mcpuser/colabfit/models/cu_mace/cu_mace__MO_000000000000_000 to calculate the energy and forces on bulk copper in FCC structure.

Relax an FCC aluminum structure with my trained model and report the final energy and cell parameters.

Generate a Python snippet to run the energy calculation on bulk silicon using my KLAY model.

Validate with test drivers:

What test drivers are available for validating my model?

Run the ElasticConstantsCrystal test driver on my silicon model at /home/mcpuser/colabfit/models/si_mace/si_mace__MO_000000000000_000.

Run the EquilibriumCrystalStructure and VacancyFormationEnergyRelaxationVolumeCrystal tests on my copper FCC model.

Check status:

Check my GPU status and list all the models and datasets I have locally.

End-to-end workflow:

Search ColabFit for silicon datasets with forces, download the best one, train a MACE model, calculate energy and forces on bulk diamond-cubic silicon, then run the ElasticConstantsCrystal and EquilibriumCrystalStructure test drivers to validate the model. Report the elastic constants and equilibrium lattice parameter when done.

Stopping / Canceling Training

The MCP server runs via docker compose run (not docker compose up), so docker compose down alone will not stop an active training container. Use the methods below to stop the server including any in-progress training job.

Using Makefile

make stop

Without 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

If the project directory is not named colabfit-mcp, replace the filter value with your directory name (lowercased). You can check the label on a running container with:

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

Training progress is saved as training.log inside the model's KIM subdirectory

( __MO_000000000000_000/training.log). Stopping mid-training discards any

in-progress epoch; completed epochs and their checkpoints are preserved on disk.

Monitoring Training Progress

View training output in the following ways:

1. Real-time Container Logs (Recommended)

View live training output as it happens:

# Using Makefile

make logs

Or directly with docker compose

docker compose logs -f server

Press Ctrl+C to exit (training continues in background).

2. Persistent Log Files

Training writes log files inside the model's KIM subdirectory:

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

GPU Support

start.sh automatically detects your GPU:

• NVIDIA GPU present: starts with compose.nvidia.yaml overlay, enabling CUDA passthrough via nvidia-container-toolkit
• No NVIDIA GPU: starts without the overlay; the container selects the best available device (MPS or CPU) automatically at runtime

The pip-installed version handles GPU detection purely in Python via detect_device() — no shell wrapper needed, since PyTorch can see the host GPU directly.

Local Installation (without Docker)

Install

pip install colabfit-mcp

This enables search_datasets, check_local_datasets, download_dataset,