# ML Training Setup MCP Server Specification

## Overview

The ML Training Setup MCP Server is a Model Context Protocol implementation that provides sequential, failure-tolerant generation of machine learning training scripts. It employs a step-by-step approach to ensure robust script generation with strict file locking constraints.

## Core Concepts

### Sequential Thinking Architecture

The server implements a sequential execution pipeline where each step:
1. Builds upon the previous successful step
2. Validates its output before proceeding
3. Can apply fixes to errors on failure

### State Management

The server maintains a persistent state that tracks:
- Current execution step
- Locked file path
- Generated artifacts
- Validation results
- Error recovery attempts

## Architecture

### Components

```
┌─────────────────────────────────────────────────┐
│                 AI Agent (Claude)                │
└─────────────────┬───────────────────────────────┘
                  │
┌─────────────────▼───────────────────────────────┐
│           ML Training Setup MCP Server          │
│                                                  │
│  ┌───────────────────────────────────────────┐  │
│  │      AI-Driven Regeneration Engine         │  │
│  │  ┌─────────┐  ┌──────────┐  ┌──────────┐ │  │
│  │  │ Analyze │→ │ Context  │→ │Regenerate│ │  │
│  │  │ Error   │  │ Provider │  │   File   │ │  │
│  │  └─────────┘  └──────────┘  └──────────┘ │  │
│  └───────────────────────────────────────────┘  │
│                                                  │
│  ┌───────────────────────────────────────────┐  │
│  │           State Manager                    │  │
│  │  • Managed Files (Training + Configs)      │  │
│  │  • Execution Tracking                      │  │
│  │  • Backup Management                       │  │
│  │  • Training Status Monitor                 │  │
│  └───────────────────────────────────────────┘  │
│                                                  │
│  ┌───────────────────────────────────────────┐  │
│  │         Execution Interface                │  │
│  │  • Managed Execution (deprecated)          │  │
│  │  • Status Tracking (needs fix)             │  │
│  │  • Hang Detection                          │  │
│  └───────────────────────────────────────────┘  │
└──────────────────────────────────────────────────┘

Data Flow for AI Regeneration:
┌──────────┐     ┌──────────┐     ┌──────────┐
│  Error   │────▶│  Context │────▶│    AI    │
│ Detected │     │ Analysis │     │ Generate │
└──────────┘     └──────────┘     └──────────┘
                                        │
                                        ▼
┌──────────┐     ┌──────────┐     ┌──────────┐
│  Fresh   │◀────│  Reset   │◀────│  Apply   │
│  Start   │     │  State   │     │   Fix    │
└──────────┘     └──────────┘     └──────────┘
```

## Tools

### 1. initialize_training_file

**AI Agent Instructions:** When this tool is called, you (the AI agent) should first analyze the provided reference (GitHub URL, HuggingFace model, etc.) to understand the training approach, then generate a complete training script based on that analysis. Pass the generated script as the 'content' parameter.

This tool saves your generated training script to a file and locks the session to work exclusively with this file.

**Input Schema:**
```json
{
  "type": "object",
  "properties": {
    "file_name": {
      "type": "string",
      "description": "Name of the training script file to create (e.g., 'train_model.py')"
    },
    "content": {
      "type": "string",
      "description": "The training script content generated by the AI agent based on the reference"
    },
    "reference": {
      "type": "string",
      "description": "The original reference (GitHub URL, HuggingFace model) used to generate this script - stored for tracking"
    }
  },
  "required": ["file_name", "content", "reference"]
}
```

**Output:**
```json
{
  "status": "success|failure",
  "file_path": "string",
  "locked_file": "string",
  "message": "string"
}
```

### 2. execute_training

**AI Agent Instructions:** DO NOT use this tool. Instead, use your native Bash tool to execute training commands directly. This provides better control over execution, environment activation, and error handling.

**Recommended approach:**
1. Use `get_current_file` to know which file to execute
2. Use your Bash tool to run commands like:
   - `uv run python train_model.py`
   - `source .venv/bin/activate && python train_model.py`
   - `accelerate launch train_model.py`

**Legacy Input Schema (kept for compatibility):**
```json
{
  "type": "object",
  "properties": {
    "command": {
      "type": "string",
      "description": "Command to execute - but prefer using Bash tool directly"
    },
    "use_locked_file": {
      "type": "boolean",
      "description": "If true, automatically appends the locked file path to the command",
      "default": true
    }
  },
  "required": ["command"]
}
```

**Output:**
```json
{
  "status": "success|failure|running",
  "execution_id": "string",
  "output": "string",
  "errors": "string",
  "exit_code": "number"
}
```

### 3. monitor_and_fix (AI Context Provider)

**AI Agent Instructions:** This tool analyzes errors and returns context for you to regenerate the code. It does NOT fix files directly. After calling this, analyze the returned context and use `regenerate_file_with_fixes` to apply corrections.

Analyzes execution errors and returns structured context for AI-driven regeneration.

**Input Schema:**
```json
{
  "type": "object",
  "properties": {
    "execution_id": {
      "type": "string",
      "description": "ID of the execution to monitor"
    },
    "error_trace": {
      "type": "string",
      "description": "Error trace to analyze"
    },
    "fix_instructions": {
      "type": "string",
      "description": "Optional hints for fixing the error"
    },
    "file_name": {
      "type": "string",
      "description": "Specific file to analyze (optional, defaults to training file)"
    }
  },
  "required": ["error_trace"]
}
```

**Output:**
```json
{
  "status": "needs_ai_regeneration",
  "current_code": "string (full file content)",
  "error_trace": "string",
  "error_analysis": {
    "type": "ModuleNotFoundError|NameError|SyntaxError|TypeError|AttributeError",
    "details": {},
    "line_number": null,
    "suggestion": "string"
  },
  "file_path": "string",
  "file_name": "string",
  "message": "AI agent should analyze this context and regenerate the corrected code"
}
```

### 4. get_current_file

Returns the content of the currently locked file.

**Input Schema:**
```json
{
  "type": "object",
  "properties": {}
}
```

**Output:**
```json
{
  "file_path": "string",
  "content": "string"
}
```

### 5. regenerate_file_with_fixes

**AI Agent Instructions:** After analyzing error context from `monitor_and_fix`, generate the complete corrected code and use this tool to replace the file. The file will be backed up automatically.

Replaces a managed file with AI-corrected version.

**Input Schema:**
```json
{
  "type": "object",
  "properties": {
    "file_name": {
      "type": "string",
      "description": "Name of the file to regenerate (e.g., 'train.py')"
    },
    "new_content": {
      "type": "string",
      "description": "Complete regenerated content with fixes applied"
    },
    "backup": {
      "type": "boolean",
      "description": "Whether to backup original file (default: true)"
    }
  },
  "required": ["file_name", "new_content"]
}
```

**Output:**
```json
{
  "status": "success|failure",
  "message": "string",
  "file_path": "string",
  "file_name": "string",
  "backup_path": "string (if backup=true)",
  "note": "Training state reset for fresh execution"
}
```

### 6. update_hyperparameters

**AI Agent Instructions:** Use this when training is taking too long and you need to adjust parameters for faster feedback. The tool returns current code and requested changes for you to regenerate with new parameters.

Requests hyperparameter updates for faster training feedback.

**Input Schema:**
```json
{
  "type": "object",
  "properties": {
    "updates": {
      "type": "object",
      "description": "Dictionary of parameter updates (e.g., {'epochs': 2, 'batch_size': 8})"
    },
    "file_name": {
      "type": "string",
      "description": "File to update (optional, defaults to training file)"
    }
  },
  "required": ["updates"]
}
```

**Output:**
```json
{
  "status": "needs_ai_update",
  "current_content": "string",
  "requested_updates": {},
  "file_path": "string",
  "file_name": "string",
  "message": "AI agent should regenerate with updated hyperparameters",
  "suggestions": ["array of suggestions for faster feedback"]
}
```

### 7. check_training_status

Timer with hang detection that reports training status. AI agents use this to decide between hyperparameter adjustment vs investigating hangs.

**Input Schema:**
```json
{
  "type": "object",
  "properties": {
    "interval_minutes": {
      "type": "number",
      "description": "Check interval in minutes (default: 5)"
    }
  }
}
```

**Output:**
```json
{
  "status": "still_running|possibly_hung|completed|failed|not_started",
  "elapsed_time_minutes": 10.5,
  "last_output_minutes_ago": 8.2,
  "interval_minutes": 5,
  "intervals_passed": 2,
  "suggestion_type": "reduce_hyperparameters|check_for_hang|none",
  "message": "Training has been running for 10.5 minutes (2 x 5 minute intervals passed), no output for 8.2 minutes",
  "training_start_time": "2024-01-01T10:00:00"
}
```

**Suggestion Types:**
- `reduce_hyperparameters`: Running > 5 mins, still producing output → AI should reduce params
- `check_for_hang`: No output for > 3 mins → AI should check GPU/process state
- `none`: Everything normal


## AI-Driven Regeneration Workflow

### Overview

The server implements an AI-driven code regeneration system that replaces traditional patching with complete file regeneration. This ensures clean, working code without syntax corruption.

### Workflow Pattern

```
1. Error Detection → check_training_status
   ↓
2. Error Analysis → monitor_and_fix (returns context)
   ↓
3. AI Regeneration → AI analyzes context and regenerates code
   ↓
4. Apply Fix → regenerate_file_with_fixes (replaces file)
   ↓
5. Fresh Start → Training state reset, retry execution
```

### Key Principles

1. **Context Over Patching**: Instead of applying patches, the system returns full context for AI analysis
2. **Complete Regeneration**: AI regenerates the entire file with fixes, ensuring consistency
3. **Automatic Backups**: Original files are backed up before regeneration
4. **State Reset**: Training state is cleared after regeneration for fresh execution

### Error Analysis Types

The `monitor_and_fix` tool identifies and categorizes errors:
- **ModuleNotFoundError**: Missing imports or uninstalled packages
- **NameError**: Undefined variables or functions
- **SyntaxError**: Code structure issues
- **TypeError**: Argument mismatches or type issues
- **AttributeError**: Missing attributes or methods

## Sequential Pipeline Steps

### Step 1: Initial File Creation
**Requirements:**
- User provides to AI agent:
  - File name for the training script
  - Reference (GitHub URL, HuggingFace model name, or model card URL)
- AI agent analyzes reference and creates training script content
- MCP server saves and locks to this file for all subsequent operations
- **File Lock:** Once created, MCP works ONLY with this file for all subsequent steps

**Process (AI Agent):**
1. Receive user requirements (file name and reference)
2. Fetch and analyze reference from:
   - GitHub repository (clone and examine training code)
   - HuggingFace model card (extract training details)
   - HuggingFace model name (fetch associated training information)
3. Generate complete training script based on reference analysis
4. Call initialize_training_file with:
   - file_name: as provided by user
   - content: the generated training script
   - reference: the original reference URL/name for tracking

**Process (MCP Server):**
1. Receive generated script from AI agent
2. Save content to specified file
3. Lock session to this file
4. Store reference for future tracking
5. Return success status with locked file path

### Step 2: Training Execution
**Purpose:** Execute the training script using AI agent's Bash tool
- **IMPORTANT:** AI agent should use their native Bash tool to execute commands, NOT the MCP execute_training tool
- AI agent determines the best execution command based on:
  - The locked file's requirements
  - Available environment (.venv at workspace root)
  - Training configuration needs
- Preferred execution patterns (for AI agent to run via Bash):
  - With venv: `uv run python train_model.py` (when .venv exists)
  - Accelerate: `uv run accelerate launch train_model.py`
  - Distributed: `uv run torchrun --nproc_per_node=4 train_model.py`
  - Fallback: `source .venv/bin/activate && python train_model.py`
- Additional command examples:
  - `CUDA_VISIBLE_DEVICES=0,1 uv run python train_model.py`
  - `uv run python -m torch.distributed.launch train_model.py`
- AI agent captures output/errors using their Bash tool
- AI agent monitors execution and handles errors

### Step 3: Error Resolution and Monitoring
**Purpose:** AI agent monitors training and adjusts for faster feedback
- **Time-based adjustments:**
  - AI agent periodically calls `check_training_status`
  - MCP returns status and suggestion_type:
    - `reduce_hyperparameters`: Training slow but working → AI reduces params
    - `check_for_hang`: No output for 3+ mins → AI checks GPU/process state
  - AI agent takes appropriate action:
    - For slow training: Modify batch_size, num_steps, sequence_length
    - For hangs: Check nvidia-smi, process state, or kill/restart
- **Error handling:**
  - If errors found in Bash output:
    - AI agent uses `monitor_and_fix` tool to fix errors
    - AI agent re-executes using Bash tool
- **Success scaling:**
  - If training completes quickly with reduced params:
    - AI agent can scale up for production run
- **Key:** MCP is just a timer, AI agent makes all intelligent decisions

### Step 4: Testing
- Run tests on the locked training file to validate functionality

### Future Steps (TODO)
- Step 5: Reference persistence and tracking
  - Store user-provided references (GitHub URLs, HuggingFace models) in state
  - Maintain reference history across sessions
  - Enable re-fetching and updating from original sources
  - Add reset/unlock tool to remove file lock constraint and start fresh
- Step 6: Model architecture refinement
- Step 7: Hyperparameter optimization  
- Step 8: Training loop enhancement
- Step 9: Distributed training setup
- Step 10: Monitoring and logging integration
- Step 11: Deployment preparation
- Step 12: Model evaluation and benchmarking
- Step 13: Export and serialization
- Step 14: Production deployment readiness

## Known Issues and Future Improvements

### Critical Issues

#### 1. Status Tracking Disconnect
**Problem**: `check_training_status` only works with deprecated `execute_training` tool. When using recommended native Bash execution, status tracking fails.

**Impact**:
- Training start time not tracked with native Bash
- Hang detection unavailable
- Progress monitoring broken

**Proposed Solution**:
```json
// New tool: start_training_session
{
  "name": "start_training_session",
  "description": "Mark training as started for external execution tracking",
  "input": {
    "session_name": "string",
    "command": "string (optional)"
  }
}

// New tool: update_training_progress
{
  "name": "update_training_progress",
  "description": "Update progress for external execution",
  "input": {
    "session_id": "string",
    "status": "running|completed|failed",
    "output": "string (optional)"
  }
}
```

#### 2. Execution Strategy Confusion
**Problem**: Tool descriptions recommend native Bash but functionality requires deprecated tool.

**Solution**: Clarify execution strategies:
- **Managed Execution** (`execute_training`): Automatic tracking, limited flexibility
- **Native Execution** (Bash tool): Full control, manual tracking required

### Improvements Needed

1. **Code Cleanup**
   - Remove unused imports: `threading`, `queue`, `sys`
   - Remove backward compatibility code (`self.locked_file`)

2. **Enhanced Error Parsing**
   - Add ML-specific error patterns
   - Better multi-line trace handling
   - Framework-specific error detection (PyTorch, TensorFlow, JAX)

3. **Configuration System**
   - Make timing thresholds configurable
   - Add environment variable support
   - User-defined hang detection settings

4. **Logging Integration**
   - Add proper logging instead of in-memory storage
   - Support for external log aggregation
   - Training metrics export

## Complete Workflow Examples

### Example 1: Managed Execution with Auto-Tracking

```python
# Step 1: Initialize training file
ai_agent.call_tool("initialize_training_file", {
    "file_name": "train.py",
    "content": generated_script,
    "reference": "https://github.com/..."
})

# Step 2: Execute with tracking (using deprecated but functional tool)
ai_agent.call_tool("execute_training", {
    "command": "python",
    "use_locked_file": true
})

# Step 3: Check status (works because we used execute_training)
status = ai_agent.call_tool("check_training_status", {
    "interval_minutes": 5
})
# Returns: elapsed time, hang detection, suggestions

# Step 4: If error, analyze and fix
if status["status"] == "failed":
    context = ai_agent.call_tool("monitor_and_fix", {
        "error_trace": error_output
    })
    
    # Step 5: AI regenerates code
    fixed_code = ai_agent.analyze_and_fix(context)
    
    # Step 6: Apply fix
    ai_agent.call_tool("regenerate_file_with_fixes", {
        "file_name": "train.py",
        "new_content": fixed_code
    })
```

### Example 2: Native Bash with Manual Tracking (Current Limitation)

```python
# Step 1: Initialize file
ai_agent.call_tool("initialize_training_file", {...})

# Step 2: Execute with native Bash
result = ai_agent.bash("uv run python train.py")

# Step 3: Check status - BROKEN!
status = ai_agent.call_tool("check_training_status", {})
# Returns: "not_started" because training_start_time not set

# Workaround: Manual error handling
if "error" in result:
    context = ai_agent.call_tool("monitor_and_fix", {
        "error_trace": result["error"]
    })
    # Continue with regeneration flow...
```

### Example 3: Hyperparameter Adjustment Flow

```python
# Training taking too long
status = ai_agent.call_tool("check_training_status", {})
if status["intervals_passed"] > 2:
    # Request hyperparameter updates
    context = ai_agent.call_tool("update_hyperparameters", {
        "updates": {
            "epochs": 2,
            "batch_size": 8,
            "max_steps": 50
        }
    })
    
    # AI regenerates with new parameters
    updated_code = ai_agent.apply_hyperparameters(context)
    
    # Apply updates
    ai_agent.call_tool("regenerate_file_with_fixes", {
        "file_name": "train.py",
        "new_content": updated_code
    })
```