Pytorch Lightning MCP

🚀 PyTorch Lightning MCP Server

PyTorch Lightning MCP Server 是一个轻量级的集成层，它通过结构化、机器可读的 API 为工具、代理和编排系统提供 PyTorch Lightning 功能。

🚀 快速开始

本项目允许你通过结构化、机器可读的 API 来使用 PyTorch Lightning，可用于工具、代理和编排系统。

✨ 主要特性

• 提供用于训练、检查、验证、测试、预测和模型检查点的结构化 API。
• 支持 PyTorch Lightning 执行。
• 具备标准输入输出和 HTTP 服务器模式。

📦 安装指南

环境要求

• Python 3.10 - 3.12
• PyTorch Lightning（兼容版本）
• uv（推荐用于依赖管理）

安装命令

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

💻 使用示例

CLI 使用

你可以通过命令行界面（CLI）运行 MCP 服务器：

# 标准输入输出服务器（默认）
uv run lightning-mcp

# HTTP 服务器
uv run lightning-mcp --http --host 0.0.0.0 --port 3333

标准输入输出示例

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

HTTP 示例

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

📚 详细文档

可用工具

MCP 服务器提供以下工具（方法）：

lightning.train

使用显式配置训练 PyTorch Lightning 模型。 输入模式：

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

lightning.inspect

检查模型或运行时环境。 输入模式：

{
  "what": "model | environment | summary",
  "model": {"_target_": "string", ...} // 模型检查时必需
}

lightning.validate

验证 PyTorch Lightning 模型。 输入模式：

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

lightning.test

测试 PyTorch Lightning 模型。 输入模式：

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

lightning.predict

使用 PyTorch Lightning 模型进行预测/推理。 输入模式：

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

lightning.checkpoint

管理模型检查点：保存、加载或列出。 输入模式：

{
  "action": "save | load | list",
  "path": "string",         // 保存/加载时使用
  "directory": "string",    // 列出时使用
  "model": { ... }           // 保存/加载时使用
}

工具发现

要在运行时列出所有可用工具及其模式，请使用以下命令：

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

测试

uv run pytest

贡献代码

请参阅 CONTRIBUTING.md 和 DEVELOPMENT.md。

📄 许可证

本项目采用 Apache 2.0 许可证。