PyTorch Lightning MCP (Model Context Protocol)

An integration layer that exposes PyTorch Lightning through a structured, machine-readable API.

Intended for programmatic use by tools, agents, and orchestration systems.

Features

• Structured training and inspection APIs
• Real PyTorch Lightning execution
• Explicit, config-driven behavior
• Safe model instantiation
• Stdio and HTTP servers
• Fully tested core logic
• Clean separation between protocol, capabilities, and transport

Project Structure

src/lightning_mcp/
├── protocol.py          # Request / response schema
├── handlers/
│   ├── train.py         # Training capability
│   └── inspect.py       # Inspection capability
├── lightning/
│   └── trainer.py       # Lightning integration boundary
├── server.py            # Stdio server
├── http_server.py       # HTTP server (FastAPI)
├── models/
│   └── simple.py        # Example LightningModule
├── tools.py             # Expose tools
tests/                   # Simple test suite

Requirements

• Python 3.10 – 3.12
• PyTorch Lightning (compatible versions)
• uv (recommended)

Installation (using uv)

1. Install uv (if not already installed)

curl -Ls https://astral.sh/uv/install.sh | sh

Restart your shell after installation.

Verify:

uv --version

2. Clone the repository

git clone https://github.com/<your-org>/lightning-mcp.git
cd lightning-mcp

3. Install dependencies

To install all dependencies (including server extras):

uv sync --all-extras

This will:

• create a local virtual environment
• install PyTorch Lightning and dependencies
• install HTTP server dependencies (FastAPI, Uvicorn)

No manual venv management is required.

Usage

Training (in-process)

from lightning_mcp.handlers.train import TrainHandler
from lightning_mcp.protocol import MCPRequest

handler = TrainHandler()

request = MCPRequest(
    id="train-1",
    method="lightning.train",
    params={
        "model": {
            "_target_": "lightning_mcp.models.simple.SimpleClassifier",
            "input_dim": 4,
            "num_classes": 3,
        },
        "trainer": {
            "max_epochs": 1,
            "accelerator": "cpu",
        },
    },
)

response = handler.handle(request)

Inspection (in-process)

from lightning_mcp.handlers.inspect import InspectHandler
from lightning_mcp.protocol import MCPRequest

handler = InspectHandler()

request = MCPRequest(
    id="inspect-1",
    method="lightning.inspect",
    params={
        "what": "environment"
    },
)

response = handler.handle(request)

Stdio Server

The stdio server reads one JSON request per line from stdin and writes one JSON response per line to stdout.

Run

uv run python -m lightning_mcp.server

Example

echo '{"id":"1","method":"lightning.inspect","params":{"what":"environment"}}' \
| uv run python -m lightning_mcp.server

HTTP Server

The HTTP server exposes a single MCP endpoint.

Run

uv run uvicorn lightning_mcp.http_server:app --host 0.0.0.0 --port 3333

Endpoint

POST /mcp

Example (curl)

curl -X POST http://localhost:3333/mcp \
  -H "Content-Type: application/json" \
  -d '{
    "id": "train-http-1",
    "method": "lightning.train",
    "params": {
      "model": {
        "_target_": "lightning_mcp.models.simple.SimpleClassifier",
        "input_dim": 4,
        "num_classes": 3
      },
      "trainer": {
        "max_epochs": 1,
        "accelerator": "cpu"
      }
    }
  }'

Testing

Run the full test suite:

uv run pytest

Docker (recommended)

You can run the MCP server using Docker

{
  "mcpServers":{
    "Lightning":{
      "command":"docker",
      "args":[
        "run",
        "--rm",
        "-i",
        "lightning-mcp:latest"
      ]
    }
  }
}

MCP Tools

The MCP server exposes a small, explicit set of tools that agents can discover and invoke dynamically.

Tool discovery is available via the standard MCP method:

{
  "method": "tools/list"
}

This returns a machine-readable description of all supported tools, including their input schemas.

Available Tools

lightning.train

Train a PyTorch Lightning model using an explicit configuration.

This tool allows an agent to:

• instantiate a LightningModule
• configure a Trainer
• execute training in-process
• receive structured training metadata

Input schema (simplified):

{
  "model": {
    "_target_": "string",
    "...": "model-specific arguments"
  },
  "trainer": {
    "...": "trainer configuration (optional)"
  }
}

The model field is required and must reference a valid LightningModule class.

lightning.inspect

Inspect a model or the runtime environment without performing training.

This tool can be used to:

• inspect model architecture and parameter counts
• inspect the execution environment (Python, Torch, Lightning versions, device availability)

Input schema (simplified):

{
  "what": "model | environment",
  "model": {
    "_target_": "string",
    "...": "model-specific arguments (required for model inspection)"
  }
}

Tool Discovery Example

Using MCP stdio:

echo '{"id":"1","method":"tools/list","params":{}}' \
| uv run python -m lightning_mcp.server

Using Docker:

echo '{"id":"1","method":"tools/list","params":{}}' \
| docker run --rm -i lightning-mcp:latest

The response contains a list of tools with their names, descriptions, and input schemas.

Demo

Below is a quick example of an agent-driven interaction with the MCP server using OpenAI tool calling. The agent decides which MCP tool to use, executes exactly one action, and then summarizes the result.

Running the demo

Start the MCP server:

uvicorn lightning_mcp.http_server:app --port 8000

INFO:     Started server process [15724]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)

Run the demo on an agent of choice For example: running using an OpenAI-based agent that adapts MCP for use with OpenAI’s tool-calling interface, we may get:

Discovered MCP tools:
  OpenAI: lightning_train  ->  MCP: lightning.train
  OpenAI: lightning_inspect  ->  MCP: lightning.inspect

Agent → MCP: lightning.inspect
Args:
{
  "what": "environment"
}

MCP Result:
{
  "python": "3.11.14 (main, Oct 31 2025, 23:15:22) [Clang 21.1.4 ]",
  "torch": "2.9.1",
  "lightning": "2.6.0",
  "cuda_available": false,
  "mps_available": true
}

The inspection of the environment reveals the following setup:

- Python version: 3.11.14
- PyTorch version: 2.9.1
- PyTorch Lightning version: 2.6.0
- CUDA support: Not available
- Apple MPS (Metal Performance Shaders) support: Available

This setup indicates that the system is equipped for machine learning tasks using PyTorch and PyTorch Lightning on Apple hardware with MPS support for accelerated computing, but without CUDA support.

Contributing

We welcome contributions! Please see CONTRIBUTING.md for guidelines.

Quick Start for Developers

1. Read the guides:

   • CONTRIBUTING.md - Contribution guidelines
   • DEVELOPMENT.md - Development setup and workflow

2. Set up development environment:

   git clone https://github.com/yourusername/pytorch-lightning-mcp.git
   cd pytorch-lightning-mcp
   uv sync --all-extras
   

3. Run tests:

   pytest tests/ -v
   

4. Check code quality:

   ruff check src/ tests/
   mypy src/lightning_mcp --ignore-missing-imports
   

5. Install pre-commit hooks:

   pre-commit install
   

Testing

• All code must have tests
• Run tests with: pytest tests/ -v
• Check coverage: pytest tests/ --cov=src/lightning_mcp
• Use tox for multi-version testing: tox

Code Standards

• Linting: Ruff
• Type checking: mypy
• Format: Black-compatible (via Ruff)
• Pre-commit hooks: Available (see DEVELOPMENT.md)

Code of Conduct

Please review our Code of Conduct.

License

This project is licensed under the Apache License 2.0.