Spark MCP (Model Context Protocol) Optimizer

Spark MCP（Model Context Protocol）优化器

本项目实现了一个用于优化 Apache Spark 代码的 Model Context Protocol（MCP）服务和客户端。该系统采用客户端-服务端架构，通过智能代码优化建议和性能分析提供支持。

工作原理

代码优化工作流

graph TB
    subgraph Input
        A[Input PySpark Code] --> |spark_code_input.py| B[run_client.py]
    end

    subgraph MCP Client
        B --> |Async HTTP| C[SparkMCPClient]
        C --> |Protocol Handler| D[Tools Interface]
    end

    subgraph MCP Server
        E[run_server.py] --> F[SparkMCPServer]
        F --> |Tool Registry| G[optimize_spark_code]
        F --> |Tool Registry| H[analyze_performance]
        F --> |Protocol Handler| I[Claude AI Integration]
    end

    subgraph Resources
        I --> |Code Analysis| J[Claude AI Model]
        J --> |Optimization| K[Optimized Code Generation]
        K --> |Validation| L[PySpark Runtime]
    end

    subgraph Output
        M[optimized_spark_code.py]
        N[performance_analysis.md]
    end

    D --> |MCP Request| F
    G --> |Generate| M
    H --> |Generate| N

    classDef client fill:#e1f5fe,stroke:#01579b
    classDef server fill:#f3e5f5,stroke:#4a148c
    classDef resource fill:#e8f5e9,stroke:#1b5e20
    classDef output fill:#fff3e0,stroke:#e65100

    class A,B,C,D client
    class E,F,G,H,I server
    class J,K,L resource
    class M,N,O output

组件详情

1. 输入层

   • spark_code_input.py：待优化的 PySpark 源代码
   • run_client.py：客户端启动和配置

2. MCP 客户端层

   • 工具接口：符合协议的工具调用

3. MCP 服务层

   • run_server.py：服务初始化
   • 工具注册表：优化和分析工具
   • 协议处理器：MCP 请求/响应管理

4. 资源层

   • Claude AI：代码分析和优化
   • PySpark 运行时：代码执行和验证

5. 输出层

   • optimized_spark_code.py：优化后的代码
   • performance_analysis.md：详细分析

此工作流说明了：

1. 提交输入 PySpark 代码
2. MCP 协议处理和路由
3. Claude AI 分析和优化
4. 代码转换和验证
5. 性能分析和报告

架构

本项目遵循 Model Context Protocol 架构，实现标准化的 AI 模型交互：

┌──────────────────┐      ┌──────────────────┐      ┌──────────────────┐
│                  │      │   MCP Server     │      │    Resources     │
│   MCP Client     │      │  (SparkMCPServer)│      │                  │
│ (SparkMCPClient) │      │                  │      │ ┌──────────────┐ │
│                  │      │    ┌─────────┐   │      │ │  Claude AI   │ │
│   ┌─────────┐    │      │    │ Tools   │   │ <──> │ │   Model      │ │
│   │ Tools   │    │      │    │Registry │   │      │ └──────────────┘ │
│   │Interface│    │ <──> │    └─────────┘   │      │                  │
│   └─────────┘    │      │    ┌─────────┐   │      │ ┌──────────────┐ │
│                  │      │    │Protocol │   │      │ │  PySpark     │ │
│                  │      │    │Handler  │   │      │ │  Runtime     │ │
│                  │      │    └─────────┘   │      │ └──────────────┘ │
└──────────────────┘      └──────────────────┘      └──────────────────┘

        │                         │                          │
        │                         │                          │
        v                         v                          v
┌──────────────┐          ┌──────────────┐           ┌──────────────┐
│  Available   │          │  Registered  │           │   External   │
│    Tools     │          │    Tools     │           │  Resources   │
├──────────────┤          ├──────────────┤           ├──────────────┤
│optimize_code │          │optimize_code │           │ Claude API   │
│analyze_perf  │          │analyze_perf  │           │ Spark Engine │
└──────────────┘          └──────────────┘           └──────────────┘

组件

1. MCP 客户端

   • 提供用于代码优化的工具接口
   • 处理与服务端的异步通信
   • 管理代码生成的文件 I/O

2. MCP 服务

   • 实现 MCP 协议处理器
   • 管理工具注册表和执行
   • 协调客户端与资源之间的交互

3. 资源

   • Claude AI：提供代码优化智能能力
   • PySpark 运行时：执行并验证优化

协议流程

1. 客户端通过 MCP 协议发送优化请求
2. 服务验证请求并调用相应工具
3. 工具利用 Claude AI 进行优化
4. 通过 MCP 响应返回优化后的代码
5. 客户端保存并验证优化后的代码

端到端功能

sequenceDiagram
    participant U as User
    participant C as MCP Client
    participant S as MCP Server
    participant AI as Claude AI
    participant P as PySpark Runtime

    U->>C: Submit Spark Code
    C->>S: Send Optimization Request
    S->>AI: Analyze Code
    AI-->>S: Optimization Suggestions
    S->>C: Return Optimized Code
    C->>P: Run Original Code
    C->>P: Run Optimized Code
    P-->>C: Execution Results
    C->>C: Generate Analysis
    C-->>U: Final Report

1. 代码提交

   • 用户将 PySpark 代码放入 v1/input/spark_code_input.py
   • MCP 客户端读取代码

2. 优化过程

   • MCP 客户端通过标准化协议连接到服务
   • 服务将代码转发给 Claude AI 进行分析
   • AI 根据最佳实践提出优化建议
   • 服务验证并处理这些建议

3. 代码生成

   • 优化后的代码保存到 v1/output/optimized_spark_code.py
   • 包含解释优化内容的详细注释
   • 在改善性能的同时保持原始代码结构

4. 性能分析

   • 在 PySpark 运行时中执行两个版本
   • 比较执行时间
   • 验证结果的正确性
   • 收集并分析指标

5. 结果生成

   • 在 v1/output/performance_analysis.md 中提供全面分析
   • 并排比较执行情况
   • 性能提升统计
   • 优化说明和理由

使用方法

要求

• Python 3.8+
• PySpark 3.2.0+
• Anthropic API Key（用于 Claude AI）

安装

pip install -r requirements.txt

快速开始

1. 将待优化的 Spark 代码添加到 input/spark_code_input.py

2. 启动 MCP 服务：

python v1/run_server.py

3. 运行客户端以优化代码：

python v1/run_client.py

这将生成两个文件：

• output/optimized_spark_example.py：优化后的 Spark 代码，包含详细的优化注释
• output/performance_analysis.md：全面的性能分析

4. 运行并比较代码版本：

python v1/run_optimized.py

这将：

• 执行原始代码和优化后的代码
• 比较执行时间和结果
• 使用执行指标更新性能分析
• 显示详细的性能提升统计

项目结构

ai-mcp/
├── input/
│   └── spark_code_input.py     # Original Spark code to optimize
├── output/
│   ├── optimized_spark_example.py  # Generated optimized code
│   └── performance_analysis.md     # Detailed performance comparison
├── spark_mcp/
│   ├── client.py               # MCP client implementation
│   └── server.py               # MCP server implementation
├── run_client.py              # Client script to optimize code
├── run_server.py              # Server startup script
└── run_optimized.py           # Script to run and compare code versions

为什么选择 MCP？

Model Context Protocol（MCP）为 Spark 代码优化带来了多项关键优势：

直接调用 Claude AI 与 MCP 服务

方面	直接调用 Claude AI	MCP 服务
集成	• 为每个团队定制集成 • 手动处理响应 • 实现重复	• 预构建的客户端库 • 自动化工作流 • 统一的接口
基础设施	• 无内置验证 • 不持久化结果 • 手动跟踪	• 自动验证 • 结果持久化 • 版本控制
上下文	• 基础代码建议 • 无执行上下文 • 优化范围有限	• 感知上下文的优化 • 完整的执行历史 • 全面改进
验证	• 需要手动测试 • 无性能指标 • 结果不确定	• 自动化测试 • 性能指标 • 已验证的结果
工作流	• 临时流程 • 未标准化 • 需要人工干预	• 结构化流程 • 标准协议 • 自动化流水线

主要区别：

1. AI 集成

方式	代码示例	优势
传统方式	client = anthropic.Client(api_key) response = client.messages.create(...)	• 设置复杂 • 自定义错误处理 • 紧密耦合
MCP	client = SparkMCPClient() result = await client.optimize_spark_code(code)	• 简单的接口 • 内置验证 • 松散耦合

2. 工具管理

方式	代码示例	优势
传统方式	class SparkOptimizer:   def register_tool(self, name, func):     self.tools[name] = func	• 手动注册 • 无验证 • 维护复杂
MCP	@register_tool("optimize_spark_code") async def optimize_spark_code(code: str):	• 自动发现 • 类型检查 • 易于扩展

3. 资源管理

方式	代码示例	优势
传统方式	def __init__(self):   self.claude = init_claude()   self.spark = init_spark()	• 手动编排 • 手动清理 • 容易出错
MCP	@requires_resources(["claude_ai", "spark"]) async def optimize_spark_code(code: str):	• 自动协调 • 生命周期管理 • 错误处理

4. 通信协议

方式	代码示例	优势
传统方式	{"type": "request",  "payload": {"code": code}}	• 自定义格式 • 手动验证 • 自定义调试
MCP	{"method": "tools/call",  "params": {"name": "optimize_code"}}	• 标准格式 • 自动验证 • 易于调试

功能

• 智能代码优化：借助 Claude AI 分析并优化 PySpark 代码
• 性能分析：详细分析原始代码与优化后代码之间的性能差异
• MCP 架构：实现模型上下文协议，以标准化与 AI 模型的交互
• 轻松集成：通过简单的客户端接口提交代码优化请求
• 代码生成：自动将优化后的代码保存到单独的文件中

高级用法

你也可以通过编程方式使用客户端：

from spark_mcp.client import SparkMCPClient

async def main():
    # Connect to the MCP server
    client = SparkMCPClient()
    await client.connect()

    # Your Spark code to optimize
    spark_code = '''
    # Your PySpark code here
    '''

    # Get optimized code with performance analysis
    optimized_code = await client.optimize_spark_code(
        code=spark_code,
        optimization_level="advanced",
        save_to_file=True  # Save to output/optimized_spark_example.py
    )
    
    # Analyze performance differences
    analysis = await client.analyze_performance(
        original_code=spark_code,
        optimized_code=optimized_code,
        save_to_file=True  # Save to output/performance_analysis.md
    )
    
    # Run both versions and compare
    # You can use the run_optimized.py script or implement your own comparison
    
    await client.close()
    
    # Analyze performance
    performance = await client.analyze_performance(spark_code, optimized_code)

    await client.close()

输入和输出示例

此代码库包含一个示例工作流：

1. 输入代码（input/spark_code_input.py）：

# Create DataFrames and join
emp_df = spark.createDataFrame(employees, ["id", "name", "age", "dept", "salary"])
dept_df = spark.createDataFrame(departments, ["dept", "location", "budget"])

# Join and analyze
result = emp_df.join(dept_df, "dept") \
    .groupBy("dept", "location") \
    .agg({"salary": "avg", "age": "avg", "id": "count"}) \
    .orderBy("dept")

2. 优化后的代码（output/optimized_spark_example.py）：

# Performance-optimized version with caching and improved configurations
spark = SparkSession.builder \
    .appName("EmployeeAnalysis") \
    .config("spark.sql.shuffle.partitions", 200) \
    .getOrCreate()

# Create and cache DataFrames
emp_df = spark.createDataFrame(employees, ["id", "name", "age", "dept", "salary"]).cache()
dept_df = spark.createDataFrame(departments, ["dept", "location", "budget"]).cache()

# Optimized join and analysis
result = emp_df.join(dept_df, "dept") \
    .groupBy("dept", "location") \
    .agg(
        avg("salary").alias("avg_salary"),
        avg("age").alias("avg_age"),
        count("id").alias("employee_count")
    ) \
    .orderBy("dept")

3. 性能分析（output/performance_analysis.md）：

## Execution Results Comparison

### Timing Comparison
- Original Code: 5.18 seconds
- Optimized Code: 0.65 seconds
- Performance Improvement: 87.4%

### Optimization Details
- Caching frequently used DataFrames
- Optimized shuffle partitions
- Improved column expressions
- Better memory management

项目结构

ai-mcp/
├── spark_mcp/
│   ├── __init__.py
│   ├── client.py      # MCP client implementation
│   └── server.py      # MCP server implementation
├── examples/
│   ├── optimize_code.py           # Example usage
│   └── optimized_spark_example.py # Generated optimized code
├── requirements.txt
└── run_server.py      # Server startup script

可用工具

1. optimize_spark_code

   • 优化 PySpark 代码，提升性能
   • 支持基础和高级优化级别
   • 自动将优化后的代码保存到 examples/optimized_spark_example.py

2. analyze_performance

   • 分析原始代码和优化后代码之间的性能差异
   • 提供以下方面的见解：
     • 性能提升
     • 资源利用率
     • 可扩展性考量
     • 潜在的权衡

环境变量

• ANTHROPIC_API_KEY：用于 Claude AI 的 Anthropic API 密钥

优化示例

系统实现了多种 PySpark 优化，包括：

• 对小表与大表的连接使用广播连接
• 高效使用窗口函数
• 策略性地缓存数据
• 优化查询计划
• 按照提升性能的原则安排操作顺序

贡献

欢迎提交问题和功能改进请求！

许可证

MIT 许可证