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, 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 for details.