npm package: @sloflash/ml-training-mcp
Description: MCP server for ML training script development with progressive scaling and intelligent recovery
Author: sloflash
Latest version: 0.1.1

# 🚀 ML Training MCP Server

An intelligent MCP server for ML training script development with progressive scaling, automatic recovery, and comprehensive testing.

## Key Features

### 🎯 Single File Focus
- **ALL** training logic stays in one `training.py` file
- No multi-file sprawl - everything in one place
- Small, manageable changes (50-100 lines max per delta)

### 📈 Progressive Scaling
- Start with minimal config (batch_size=1, seq_len=128)
- Gradually scale up to target configuration
- Automatic rollback on failure
- Memory of successful configurations

### 🔍 Suspicious Pattern Detection
- Detects loss increasing after 20 steps
- Catches NaN/Inf in gradients
- Identifies memory leaks
- Spots training stalls

### 🔄 Smart Recovery
- Analyzes failures intelligently
- Suggests specific fixes (not blind retries)
- Maximum 20 retries with 10-minute timeout
- Different fix attempted each time

### 🧪 Comprehensive Testing
- Mock scenarios for common failures
- Test before running on real GPUs
- Simulate long-running training
- Validate scaling paths safely

## Quick Install (Recommended)

```bash
# No installation needed - use npx directly in your project
cd your-ml-project

# Create Claude configuration for THIS project only
mkdir -p .claude
cat > .claude/claude_desktop_config.json << 'EOF'
{
  "mcpServers": {
    "ml-training": {
      "command": "npx",
      "args": ["-y", "@sloflash/ml-training-mcp"],
      "env": {
        "PROJECT_ROOT": "${workspaceFolder}"
      }
    }
  }
}
EOF

# That's it! The MCP server is now available in THIS project only
```

## Alternative Installation Methods

### Method 1: Project-specific installation (for teams)

```bash
# Install as dev dependency in your ML project
cd your-ml-project
npm install --save-dev @sloflash/ml-training-mcp

# Configure Claude for this project
mkdir -p .claude
cat > .claude/claude_desktop_config.json << 'EOF'
{
  "mcpServers": {
    "ml-training": {
      "command": "node",
      "args": ["./node_modules/@sloflash/ml-training-mcp/dist/index.js"],
      "env": {
        "PROJECT_ROOT": "${workspaceFolder}"
      }
    }
  }
}
EOF
```

### Method 2: Global Installation (if you prefer)

```bash
# Install globally
npm install -g @sloflash/ml-training-mcp

# Then in your project, configure Claude
mkdir -p .claude
cat > .claude/claude_desktop_config.json << 'EOF'
{
  "mcpServers": {
    "ml-training": {
      "command": "ml-training-mcp",
      "args": [],
      "env": {
        "PROJECT_ROOT": "${workspaceFolder}"
      }
    }
  }
}
EOF
```

## One-Click Setup Script

Save this as `setup-mcp.sh` in your ML project:

```bash
#!/bin/bash

echo "🚀 Setting up ML Training MCP for this project..."

# Check Node.js
if ! command -v node &> /dev/null; then
    echo "❌ Node.js is required. Install from nodejs.org"
    exit 1
fi

# Create Claude config
mkdir -p .claude
cat > .claude/claude_desktop_config.json << 'EOF'
{
  "mcpServers": {
    "ml-training": {
      "command": "npx",
      "args": ["-y", "@sloflash/ml-training-mcp"],
      "env": {
        "PROJECT_ROOT": "${workspaceFolder}",
        "SINGLE_FILE_ENFORCEMENT": "true",
        "MAX_RETRIES": "20",
        "TIMEOUT_MINUTES": "10"
      }
    }
  }
}
EOF

echo "✅ MCP configured for this project only"
echo "📁 Config: .claude/claude_desktop_config.json"
echo "🔄 Restart Claude Desktop to activate"
```

Then run:
```bash
chmod +x setup-mcp.sh
./setup-mcp.sh
```

### Basic Usage

```typescript
// 1. Initialize training script
await mcp.init_training_script({ model: "meta-llama/Llama-2-7b-hf" });

// 2. Test with minimal configuration
await mcp.test_minimal_config();

// 3. Scale up progressively
await mcp.scale_up_safely();

// 4. Run training with monitoring
await mcp.run_training({ 
  timeout: 600000,  // 10 minutes
  mock: false       // Use real training
});
```

## Available Tools

### Initialization
- `init_training_script` - Create training.py from template
- `setup_environment` - Initialize uv and .venv
- `configure_model` - Set HuggingFace model/tokenizer

### Execution
- `run_training` - Execute with timeout and monitoring
- `run_mock_training` - Test with simulated scenarios
- `test_minimal_config` - Validate minimal setup

### Monitoring
- `detect_suspicious_patterns` - Check for training issues
- `get_training_status` - Current status and retries
- `monitor_memory` - Track GPU memory usage

### Recovery
- `analyze_failure` - Determine failure cause
- `suggest_fix` - Get intelligent fix suggestions
- `rollback_to_working` - Revert to last good version
- `apply_small_delta` - Make controlled changes

### Scaling
- `scale_up_safely` - Increase parameters gradually
- `scale_down` - Reduce after failure
- `find_max_viable` - Determine hardware limits

## Testing

### Run All Tests
```bash
npm test
```

### Specific Test Suites
```bash
npm run test:suspicious-patterns  # Pattern detection
npm run test:scaling              # Progressive scaling
npm run test:recovery             # Failure recovery
```

### Integration Testing
```bash
npm run test:integration
```

## Mock Scenarios

The server includes predefined scenarios for testing:

1. **happy_path** - Everything works perfectly
2. **oom_after_warmup** - OOM after 100 steps
3. **gradient_explosion** - NaN loss at step 250
4. **memory_leak** - Gradual memory growth
5. **loss_increase** - Loss going up instead of down
6. **nccl_timeout** - Multi-GPU communication failure

## Example Workflow

```typescript
// Start with mock testing
const mockResult = await mcp.run_mock_training({
  scenario: 'oom_after_warmup',
  config: { batch_size: 4, seq_len: 512 }
});

if (!mockResult.success) {
  // Analyze the failure
  const analysis = await mcp.analyze_failure({ 
    error: mockResult.error 
  });
  
  // Get suggested fix
  const fix = await mcp.suggest_fix();
  
  // Apply the fix
  await mcp.apply_small_delta({
    changes: fix.code,
    description: fix.description
  });
  
  // Retry with smaller config
  await mcp.scale_down({ factor: 0.5 });
}
```

## Suspicious Pattern Handling

When the system detects issues:

| Pattern | Detection | Automatic Action |
|---------|-----------|------------------|
| Loss increasing | After 20 steps | Reduce LR by 10x |
| Loss NaN | Immediate | Stop training |
| Memory leak | Continuous growth | Clear cache |
| No progress | 100 steps stalled | Check gradients |
| Gradient explosion | norm > 100 | Add clipping |

## Configuration

Default configuration in `Config` class:
```python
batch_size = 1
seq_len = 128  
learning_rate = 1e-5
gradient_clip = 1.0
max_steps = 1000
```

Override via command line:
```bash
uv run python training.py --batch_size 4 --seq_len 512
```

## File Structure

```
ml-training-mcp/
├── src/                      # TypeScript source
│   ├── index.ts             # MCP server
│   ├── script-manager.ts    # Version control
│   ├── suspicious-detector.ts # Pattern detection
│   ├── scaling-orchestrator.ts # Progressive scaling
│   ├── retry-manager.ts     # Retry logic
│   └── mock-executor.ts     # Test scenarios
├── storage/                  # Runtime storage
│   ├── training.py          # THE training script
│   ├── training_memory.json # Working memory
│   └── working_versions/    # Checkpoints
├── templates/               # Script templates
│   └── training_template.py
└── test/                    # Test suites
```

## Important Notes

1. **Single File Rule**: Everything stays in `training.py`
2. **Environment**: Always uses `uv` with `.venv`
3. **Testing First**: Always test with mocks before real training
4. **Small Changes**: Maximum 50-100 lines per modification
5. **Smart Recovery**: Analyzes failures, doesn't blind retry

## Troubleshooting

### OOM Errors
- Reduce batch_size by 50%
- Reduce sequence_length by 50%
- Enable gradient checkpointing
- Use gradient accumulation

### Loss Increasing
- Check learning rate (usually too high)
- Verify data shuffling
- Check gradient flow
- Add warmup schedule

### NaN Loss
- Add gradient clipping
- Reduce learning rate significantly
- Check model initialization
- Verify data normalization

## Contributing

See [CLAUDE.md](CLAUDE.md) for development guidelines and testing requirements.

## License

MIT