# 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 with checkpoint-based rollback capabilities to ensure robust script generation.

## Core Concepts

### Sequential Thinking Architecture

The server implements a sequential execution pipeline where each step:
1. Builds upon the previous successful step
2. Creates a checkpoint before execution
3. Validates its output before proceeding
4. Can rollback to the last successful state on failure

### State Management

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

## Architecture

### Components

```
┌─────────────────────────────────────────────────┐
│                 MCP Client                       │
└─────────────────┬───────────────────────────────┘
                  │
┌─────────────────▼───────────────────────────────┐
│           ML Training Setup MCP Server          │
│                                                  │
│  ┌───────────────────────────────────────────┐  │
│  │         Sequential Pipeline Manager        │  │
│  │  ┌─────────┐  ┌─────────┐  ┌──────────┐  │  │
│  │  │ Step 1  │→ │ Step 2  │→ │  [Next]  │  │  │
│  │  └─────────┘  └─────────┘  └──────────┘  │  │
│  └───────────────────────────────────────────┘  │
│                                                  │
│  ┌───────────────────────────────────────────┐  │
│  │           State Manager                    │  │
│  │  • Checkpoints                             │  │
│  │  • Rollback History                        │  │
│  │  • Error Recovery                          │  │
│  └───────────────────────────────────────────┘  │
│                                                  │
│  ┌───────────────────────────────────────────┐  │
│  │         Agent Interface                    │  │
│  │  • LLM Integration                         │  │
│  │  • Prompt Engineering                      │  │
│  │  • Response Validation                     │  │
│  └───────────────────────────────────────────┘  │
└──────────────────────────────────────────────────┘
```

## 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",
  "checkpoint_id": "string",
  "locked_file": "string",
  "message": "string"
}
```

### 2. execute_training (DEPRECATED - Use Bash tool instead)

**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

Monitors execution output and fixes errors by modifying the locked file.

**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 and fix"
    },
    "fix_instructions": {
      "type": "string",
      "description": "Specific instructions for fixing the error"
    }
  },
  "required": ["error_trace"]
}
```

### 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",
  "checkpoint_id": "string"
}
```


## 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 checkpoint ID for rollback capability

### 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 and fixes errors
- AI agent parses Bash execution output for errors
- If errors found:
  - AI agent analyzes the error trace from Bash output
  - AI agent uses `monitor_and_fix` tool to modify the locked file ONLY
  - AI agent re-executes using Bash tool after fixes
- AI agent monitors metrics from Bash output (loss, accuracy, etc.)
- Continue until successful completion or max retries
- **Key:** All execution happens via AI's Bash tool, MCP only handles file modifications

### 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