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

1. Clone the repository:

git clone <repository-url>
cd mlflow-mcp-server

2. 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 ID
• get_experiment_by_name(experiment_name: str) - Get experiment by name
• search_experiments(name?: str, token?: str) - Search experiments with optional filtering

Run Tools

• get_run(run_id: str) - Get detailed run information
• get_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 models
• get_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:

1. Create a new function in the appropriate tool file
2. Add the tool to server.py:

   from mlflow_mcp_server.tools.your_module import your_function
   mcp.add_tool(your_function)
   

Contributing

1. Fork the repository
2. Create a feature branch
3. Make your changes
4. Add tests if applicable
5. 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