MLflow MCP Server
A Model Context Protocol (MCP) server that provides seamless integration with MLflow, enabling AI assistants to interact with MLflow experiments, runs, and registered models.
Overview
This MCP server exposes MLflow functionality through a standardized protocol, allowing AI assistants like Claude to:
- Browse and search MLflow experiments
- Retrieve experiment runs and their details
- Query registered models and model versions
- Access metrics, parameters, and metadata
Features
Experiment Management
- Get Experiment: Retrieve experiment details by ID
- Get Experiment by Name: Find experiments by name
- Search Experiments: List and filter experiments with pagination support
Run Management
- Get Run: Retrieve detailed run information including metrics, parameters, and tags
- Get Experiment Runs: List all runs for a specific experiment
- Run Type Detection: Automatically identifies parent, child, or standalone runs
Model Registry
- Get Registered Models: Search and list registered models
- Get Model Versions: Browse model versions with filtering capabilities
Installation
Prerequisites
- Python 3.11 or higher
- uv package manager
Setup
- Clone the repository:
git clone <repository-url>
cd mlflow-mcp-server
- Install dependencies:
uv sync
Configuration
MLflow Connection
The server is pre-configured to connect to your internal MLflow instance:
- Tracking URI:
YOUR URI
To use with a different MLflow instance, modify mlflow_mcp_server/utils/mlflow_client.py:
import mlflow
from mlflow import MlflowClient
mlflow.set_tracking_uri("your-mlflow-tracking-uri")
client = MlflowClient()
MCP Configuration
Add the following configuration to your MCP client (e.g., ~/.cursor/mcp.json for Cursor):
{
"mcpServers": {
"mlflow": {
"command": "uvx",
"args": ["--from", "git+https://github.com/yesid-lopez/mlflow-mcp-server", "mlflow_mcp_server"],
"env": {
"MLFLOW_TRACKING_URI": "YOUR_TRACKING_URI"
}
}
}
}
Replace /path/to/mlflow-mcp-server with the actual path to your project directory.
Usage
Running the Server
uv run -m mlflow_mcp_server
Available Tools
Once configured, the following tools become available to your AI assistant:
Experiment Tools
get_experiment(experiment_id: str)- Get experiment details by IDget_experiment_by_name(experiment_name: str)- Get experiment by namesearch_experiments(name?: str, token?: str)- Search experiments with optional filtering
Run Tools
get_run(run_id: str)- Get detailed run informationget_experiment_runs(experiment_id: str, token?: str)- List runs for an experiment
Model Registry Tools
get_registered_models(model_name?: str, token?: str)- Search registered modelsget_model_versions(model_name?: str, token?: str)- Browse model versions
Example Usage with AI Assistant
You can now ask your AI assistant questions like:
- "Show me all experiments containing 'recommendation' in the name"
- "Get the details of run ID abc123 including its metrics and parameters"
- "List all registered models and their latest versions"
- "Find experiments related to customer segmentation"
Development
Project Structure
mlflow-mcp-server/
├── mlflow_mcp_server/
│ ├── __main__.py # Server entry point
│ ├── server.py # Main MCP server configuration
│ ├── tools/ # MLflow integration tools
│ │ ├── experiment_tools.py
│ │ ├── run_tools.py
│ │ └── registered_models.py
│ └── utils/
│ └── mlflow_client.py # MLflow client configuration
├── pyproject.toml # Project dependencies
└── README.md
Dependencies
- mcp[cli]: Model Context Protocol framework
- mlflow: MLflow client library
- pydantic: Data validation and serialization
Adding New Tools
To add new MLflow functionality:
- Create a new function in the appropriate tool file
- Add the tool to
server.py:from mlflow_mcp_server.tools.your_module import your_function mcp.add_tool(your_function)
Contributing
- Fork the repository
- Create a feature branch
- Make your changes
- Add tests if applicable
- Submit a pull request
License
[Add your license information here]
Support
For issues and questions:
- Check existing issues in the repository
- Create a new issue with detailed reproduction steps