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 OR client-provided content (JSONL, CSV, Parquet)
- Upload datasets directly without server filesystem access
- 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
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
ml-lab init
ml-lab vault create
2. Add Provider Credentials
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:
{
"mcpServers": {
"ml-lab": {
"command": "ml-lab",
"args": ["serve"]
}
}
}
MCP Tools
Credentials
Datasets
Experiments
Training
Infrastructure
Remote VPS
Ollama
Open WebUI
Security
Codex Integration (Executor LLM)
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)
Analysis Types:
training- Analyze loss curves, convergence, hyperparametersexperiment- Compare experiments, suggest improvementsactivity- Review audit logs for patterns/anomaliescost- Cost efficiency analysis and recommendationsdataset- Data quality and training suitability
Schedule Frequencies:
after_training- Auto-analyze when training completeshourly/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 for details.