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. パフォーマンスの分析とレポート

アーキテクチャ

このプロジェクトは、AI モデルとのやり取りを標準化するため、Model Context Protocol アーキテクチャに従っています。

┌──────────────────┐      ┌──────────────────┐      ┌──────────────────┐
│                  │      │   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 クライアント

   • コード最適化のためのツールインターフェースを提供
   • サーバーとの非同期通信を処理
   • コード生成のためのファイル入出力を管理

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

次の 2 つのファイルが生成されます。

• 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 モデルとのやり取りを標準化するため、Model Context Protocol を実装しています
• 簡単な統合：コード最適化リクエスト用のシンプルなクライアントインターフェース
• コード生成：最適化したコードを別のファイルに自動保存します

高度な使い方

クライアントをプログラムから使用することもできます：

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 の最適化を実装しています：

• 小規模テーブルと大規模テーブルの結合におけるブロードキャスト結合
• 効率的なウィンドウ関数の使用
• 戦略的なデータキャッシュ
• クエリプランの最適化
• パフォーマンス指向の操作順序

コントリビューション

Issue や機能改善のリクエストをぜひお寄せください！

ライセンス

MIT License