PyTorch Lightning MCP Server

A minimal integration layer exposing PyTorch Lightning via a structured, machine-readable API for tools, agents, and orchestration systems.

Features

• Structured APIs for training, inspecting, validating, testing, predicting, and checkpointing models
• PyTorch Lightning execution
• Stdio and HTTP server modes

Requirements

• Python 3.10–3.12
• PyTorch Lightning (compatible version)
• uv (recommended for dependency management)

Installation

curl -Ls https://astral.sh/uv/install.sh | sh
git clone https://github.com/<your-org>/lightning-mcp.git
cd lightning-mcp
uv sync --all-extras

Usage

CLI

You can run the MCP server via CLI:

# Stdio server (default)
uv run lightning-mcp

# HTTP server
uv run lightning-mcp --http --host 0.0.0.0 --port 3333

Stdio Example

echo '{"id":"1","method":"lightning.inspect","params":{"what":"environment"}}' | uv run lightning-mcp

HTTP Example

curl -X POST http://localhost:3333/mcp \
  -H "Content-Type: application/json" \
  -d '{"id":"1","method":"lightning.inspect","params":{"what":"environment"}}'

Available Tools

The MCP server exposes the following tools (methods):

lightning.train

Train a PyTorch Lightning model with explicit configuration.

Input schema:

{
  "model": {"_target_": "string", ...},
  "trainer": { ... }
}

lightning.inspect

Inspect a model or the runtime environment.

Input schema:

{
  "what": "model | environment | summary",
  "model": {"_target_": "string", ...} // required for model inspection
}

lightning.validate

Validate a PyTorch Lightning model.

Input schema:

{
  "model": {"_target_": "string", ...},
  "trainer": { ... }
}

lightning.test

Test a PyTorch Lightning model.

Input schema:

{
  "model": {"_target_": "string", ...},
  "trainer": { ... }
}

lightning.predict

Run prediction/inference with a PyTorch Lightning model.

Input schema:

{
  "model": {"_target_": "string", ...},
  "trainer": { ... }
}

lightning.checkpoint

Manage model checkpoints: save, load, or list.

Input schema:

{
  "action": "save | load | list",
  "path": "string",         // for save/load
  "directory": "string",    // for list
  "model": { ... }           // for save/load
}

Tool Discovery

To list all available tools and their schemas at runtime:

echo '{"id":"1","method":"tools/list","params":{}}' | uv run lightning-mcp

Testing

uv run pytest

Contributing

See CONTRIBUTING.md and DEVELOPMENT.md.

License

Apache 2.0