# ML Lab MCP
A comprehensive MCP (Model Context Protocol) server for ML model training, fine-tuning, and experimentation. Transform your AI assistant into a full ML engineering environment.
## Features
### Unified Credential Management
- Encrypted vault for API keys (Lambda Labs, RunPod, Mistral, OpenAI, Together AI, etc.)
- PBKDF2 key derivation with AES encryption
- Never stores credentials in plaintext
### Dataset Management
- Register datasets from local files (JSONL, CSV, Parquet)
- Automatic schema inference and statistics
- Train/val/test splitting
- Template-based transformations
### Experiment Tracking
- SQLite-backed experiment storage
- Version control and comparison
- Fork experiments with config modifications
- Full metrics history
### Multi-Backend Training
- **Local**: transformers + peft + trl for local GPU training
- **Mistral API**: Native fine-tuning for Mistral models
- **Together AI**: Hosted fine-tuning service
- **OpenAI**: GPT model fine-tuning
### Cloud GPU Provisioning
- **Lambda Labs**: H100, A100 instances
- **RunPod**: Spot and on-demand GPUs
- Automatic price comparison across providers
- Smart routing based on cost and availability
### Remote VPS Support
- Use any SSH-accessible machine (Hetzner, Hostinger, OVH, home server, university cluster)
- Automatic environment setup
- Dataset sync via rsync
- Training runs in tmux (persistent across disconnects)
- Amortized hourly cost calculation from monthly fees
### Cost Estimation
- Pre-training cost estimates across all providers
- Real-time pricing queries
- Time estimates based on model and dataset size
### Ollama Integration
- Deploy fine-tuned GGUF models to Ollama
- Pull models from Ollama registry
- Chat/inference testing directly from MCP
- Model management (list, delete, copy)
### Open WebUI Integration
- Create model presets with system prompts
- Knowledge base management (RAG)
- Chat through Open WebUI (applies configs + knowledge)
- Seamless Ollama ↔ Open WebUI workflow
## Installation
```bash
pip install ml-lab-mcp
# With training dependencies
pip install ml-lab-mcp[training]
# With cloud provider support
pip install ml-lab-mcp[cloud]
# Everything
pip install ml-lab-mcp[training,cloud,dev]
```
## Quick Start
### 1. Initialize and Create Vault
```bash
ml-lab init
ml-lab vault create
```
### 2. Add Provider Credentials
```bash
ml-lab vault unlock
ml-lab vault add --provider lambda_labs --api-key YOUR_KEY
ml-lab vault add --provider mistral --api-key YOUR_KEY
```
### 3. Configure with Claude Code / Claude Desktop
Add to your MCP configuration:
```json
{
"mcpServers": {
"ml-lab": {
"command": "ml-lab",
"args": ["serve"]
}
}
}
```
## MCP Tools
### Credentials
| Tool | Description |
|------|-------------|
| `creds_create_vault` | Create encrypted credential vault |
| `creds_unlock` | Unlock vault with password |
| `creds_add` | Add provider credentials |
| `creds_list` | List configured providers |
| `creds_test` | Verify credentials work |
### Datasets
| Tool | Description |
|------|-------------|
| `dataset_register` | Register a local dataset file |
| `dataset_list` | List all datasets |
| `dataset_inspect` | View schema and statistics |
| `dataset_preview` | Preview samples |
| `dataset_split` | Create train/val/test splits |
| `dataset_transform` | Apply template transformations |
### Experiments
| Tool | Description |
|------|-------------|
| `experiment_create` | Create new experiment |
| `experiment_list` | List experiments |
| `experiment_get` | Get experiment details |
| `experiment_compare` | Compare multiple experiments |
| `experiment_fork` | Fork with modifications |
### Training
| Tool | Description |
|------|-------------|
| `train_estimate` | Estimate cost/time across providers |
| `train_launch` | Start training run |
| `train_status` | Check run status |
| `train_stop` | Stop training |
### Infrastructure
| Tool | Description |
|------|-------------|
| `infra_list_gpus` | List available GPUs with pricing |
| `infra_provision` | Provision cloud instance |
| `infra_terminate` | Terminate instance |
### Remote VPS
| Tool | Description |
|------|-------------|
| `vps_register` | Register a VPS (host, user, key, GPU info, monthly cost) |
| `vps_list` | List all registered VPS machines |
| `vps_status` | Check VPS status (online, GPU, running jobs) |
| `vps_unregister` | Remove a VPS from registry |
| `vps_setup` | Install training dependencies on VPS |
| `vps_sync` | Sync dataset to VPS |
| `vps_run` | Run command on VPS |
| `vps_logs` | Get training logs from VPS |
### Ollama
| Tool | Description |
|------|-------------|
| `ollama_status` | Check Ollama status (running, version, GPU) |
| `ollama_list` | List models in Ollama |
| `ollama_pull` | Pull model from registry |
| `ollama_deploy` | Deploy GGUF to Ollama |
| `ollama_chat` | Chat with a model |
| `ollama_delete` | Delete a model |
### Open WebUI
| Tool | Description |
|------|-------------|
| `owui_status` | Check Open WebUI connection |
| `owui_list_models` | List model configurations |
| `owui_create_model` | Create model preset (system prompt, params) |
| `owui_delete_model` | Delete model configuration |
| `owui_list_knowledge` | List knowledge bases |
| `owui_create_knowledge` | Create knowledge base |
| `owui_add_knowledge_file` | Add file to knowledge base |
| `owui_chat` | Chat through Open WebUI |
### Security
| Tool | Description |
|------|-------------|
| `security_audit_log` | View recent audit log entries |
| `security_audit_summary` | Get audit activity summary |
| `security_tailscale_status` | Check Tailscale VPN connection |
| `security_ssh_key_rotate` | Rotate SSH key for a VPS |
| `creds_expiry_check` | Check credential expiry status |
| `creds_rotate` | Rotate credentials for a provider |
### Codex Integration (Executor LLM)
| Tool | Description |
|------|-------------|
| `codex_status` | Check if Codex CLI is available |
| `codex_analyze_error` | Have Codex analyze errors and suggest fixes |
| `codex_generate_training_script` | Generate training script from experiment config |
| `codex_fix_code` | Have Codex fix issues in training code |
| `codex_optimize_config` | Optimize training config for quality/speed/memory |
| `codex_debug_training` | Debug training issues from logs |
| `codex_run` | Run arbitrary task with Codex |
**Architecture**: Separation of intelligence vs execution:
- **Planner LLM (Claude)**: Reasoning, architecture, tradeoffs
- **Executor LLM (Codex)**: Precise code edits + CLI work
- **World tools (ML Lab)**: Data, training, infra, deployment
### Deep Thinking Analysis (Ollama Reasoning Models)
| Tool | Description |
|------|-------------|
| `thinking_analyze` | Run deep analysis with reasoning models (DeepSeek R1, QwQ) |
| `thinking_schedule` | Schedule automated analysis (after training, daily, weekly) |
| `thinking_reports` | List and retrieve analysis reports |
**Analysis Types**:
- `training` - Analyze loss curves, convergence, hyperparameters
- `experiment` - Compare experiments, suggest improvements
- `activity` - Review audit logs for patterns/anomalies
- `cost` - Cost efficiency analysis and recommendations
- `dataset` - Data quality and training suitability
**Schedule Frequencies**:
- `after_training` - Auto-analyze when training completes
- `hourly` / `daily` / `weekly` - Periodic reviews
Reports stored in `~/.cache/ml-lab/reports/` for historical reference.
## Example Workflow
```
User: Fine-tune Llama 3.1 8B on my customer support data
Claude: I'll help you fine-tune Llama 3.1 8B. Let me first check your data.
[Uses dataset_register to load ./support_data.jsonl]
Found 15,000 samples. Let me estimate the costs:
[Uses train_estimate]
| Provider | Time | Cost | Availability |
|---------------|--------|--------|--------------|
| vps:hetzner-01| ~4h | $0.28* | Idle |
| Local (4090) | ~12h | $0 | Ready |
| Lambda H100 | ~2h | $5.50 | 3 available |
| Together API | ~4h | $12 | Queued |
*amortized from $200/mo
Which would you prefer?
User: Use Lambda Labs
Claude: [Uses experiment_create, then train_launch with backend=lambda]
Training started! Run ID: abc123
I'll monitor progress and let you know when it completes.
```
## Architecture
```
src/ml_lab/
├── server.py # MCP server entry point (61 tools)
├── credentials.py # Encrypted credential vault
├── cli.py # Command-line interface
├── backends/
│ ├── base.py # Training backend interface
│ ├── local.py # Local GPU training
│ ├── mistral_api.py # Mistral fine-tuning API
│ ├── together_api.py # Together AI API
│ ├── openai_api.py # OpenAI fine-tuning API
│ └── vertex_api.py # Google Vertex AI (Gemini)
├── cloud/
│ ├── base.py # Cloud provider interface
│ ├── lambda_labs.py # Lambda Labs integration
│ ├── runpod.py # RunPod integration
│ ├── modal_provider.py # Modal integration
│ └── remote_vps.py # Generic SSH VPS support (+ Tailscale)
├── storage/
│ ├── datasets.py # Dataset management
│ └── experiments.py # Experiment tracking
├── inference/
│ ├── ollama.py # Ollama integration
│ ├── openwebui.py # Open WebUI integration
│ └── thinking.py # Deep thinking analysis (DeepSeek R1, QwQ)
├── integrations/
│ └── codex.py # Codex CLI integration (executor LLM)
├── security/
│ └── audit.py # Audit logging
└── evals/
└── benchmarks.py # Evaluation suite
```
## Security
- Credentials encrypted with Fernet (AES-128-CBC)
- PBKDF2-SHA256 key derivation (480,000 iterations)
- Vault file permissions set to 600 (owner read/write only)
- API keys never logged or transmitted unencrypted
- **Audit logging**: All sensitive operations logged to `~/.cache/ml-lab/audit.log`
- **Credential expiry**: Automatic tracking with rotation reminders
- **Tailscale support**: Optional VPN requirement for VPS connections
- **SSH key rotation**: Automated rotation with rollback on failure
## Supported Providers
### Compute Providers
- Lambda Labs (H100, A100, A10)
- RunPod (H100, A100, RTX 4090)
- Modal (serverless GPU functions)
### Fine-Tuning APIs
- Mistral AI (Mistral, Mixtral, Codestral)
- Together AI (Llama, Mistral, Qwen)
- OpenAI (GPT-4o, GPT-3.5)
- Google Vertex AI (Gemini 1.5 Pro, Gemini 1.5 Flash)
### Model Hubs
- Hugging Face Hub
- Replicate
- Ollama (local GGUF models)
## Contributing
Contributions welcome! Please read CONTRIBUTING.md for guidelines.
## License
PolyForm Noncommercial 1.0.0 - free for personal use, contact for commercial licensing.
See [LICENSE](LICENSE) for details.
