Labsco
yesid-lopez logo

MLflow MCP Server

from yesid-lopez

Integrates with MLflow, enabling AI assistants to interact with experiments, runs, and registered models.

๐Ÿ”ฅ๐Ÿ”ฅ๐Ÿ”ฅโœ“ VerifiedFreeQuick setup

MLflow MCP Server

A Model Context Protocol (MCP) server that exposes MLflow experiment tracking and model registry operations as tools for AI assistants.

Table of Contents

Tools

Experiment Management

ToolDescription
get_experimentGet experiment details by ID
get_experiment_by_nameGet experiment details by name
search_experimentsList and filter experiments with optional name matching and pagination

Run Management

ToolDescription
get_runGet full run details including metrics, parameters, tags, and run type (parent/child/standalone)
get_experiment_runsList runs for an experiment with pagination

Model Registry

ToolDescription
get_registered_modelsSearch and list registered models
get_model_versionsBrowse model versions with filtering
create_registered_modelCreate a new registered model with optional description and tags
create_model_versionCreate a new model version from a run's artifacts
rename_registered_modelRename an existing registered model
set_registered_model_aliasAssign an alias (e.g. champion, challenger) to a model version
delete_registered_modelDelete a registered model and all its versions
delete_model_versionDelete a specific model version

Example Prompts

Once configured, you can ask your AI assistant things like:

Exploring experiments and runs:

  • "List all experiments related to recommendation models"
  • "Show me the runs for experiment 12 and compare their metrics"
  • "Get the parameters and metrics for run abc123"
  • "Which runs in the fraud-detection experiment have the highest accuracy?"

Managing the model registry:

  • "Show me all registered models"
  • "Register a new model called churn-classifier with description 'Binary classifier for customer churn'"
  • "Create a new version of churn-classifier from run abc123"
  • "Set the champion alias on version 3 of churn-classifier"
  • "Rename the model old-name to new-name"
  • "Delete version 1 of churn-classifier"

Analysis and comparison:

  • "Compare the last 5 runs of the search-ranking experiment by NDCG and latency"
  • "What hyperparameters were used in the best-performing run of experiment 7?"
  • "List all model versions for recommendation-model and their aliases"

Project Structure

mlflow_mcp_server/
โ”œโ”€โ”€ __main__.py              # Entry point
โ”œโ”€โ”€ server.py                # MCP server setup and tool registration
โ”œโ”€โ”€ tools/
โ”‚   โ”œโ”€โ”€ experiment_tools.py  # Experiment search and retrieval
โ”‚   โ”œโ”€โ”€ run_tools.py         # Run details and listing
โ”‚   โ””โ”€โ”€ registered_models.py # Model registry CRUD operations
โ””โ”€โ”€ utils/
    โ””โ”€โ”€ mlflow_client.py     # MLflow client singleton

Adding New Tools

  1. Create a function in the appropriate file under tools/.
  2. Register it in server.py:
from mlflow_mcp_server.tools.your_module import your_function
mcp.add_tool(your_function)

Linting

uv run ruff check .
uv run ruff format --check .