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

| 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 (Lambda Labs, GCP, OpenAI supported) |

Datasets

| Tool | Description | |------|-------------| | dataset_register | Register a dataset from a local file | | dataset_register_content | Register a dataset from client-provided content (CSV, JSON, JSONL, Parquet) | | 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 for details.