Mlflow Experiment Tracker

Model Registry Operations

• Register trained models
• Manage model versions
• Transition models between stages (Staging, Production, Archived)
• Add model descriptions and tags

Run Comparison and Analysis

• Compare metrics across runs
• Search runs by parameters/metrics
• Retrieve best performing runs
• Generate comparison visualizations

Prerequisites

MLflow Installation

pip install mlflow>=2.0.0

MLflow Tracking Server

Configure tracking URI:

import mlflow
mlflow.set_tracking_uri("http://localhost:5000")  # or remote server

Optional: MLflow MCP Server

For enhanced LLM integration, install the MLflow MCP server:

pip install mlflow>=3.4  # Official MCP support
# or
pip install mlflow-mcp   # Community server

Usage Patterns

Starting an Experiment Run

import mlflow

# Set experiment
mlflow.set_experiment("my-classification-experiment")

# Start run with context manager
with mlflow.start_run(run_name="baseline-model"):
    # Log parameters
    mlflow.log_param("learning_rate", 0.01)
    mlflow.log_param("batch_size", 32)
    mlflow.log_param("epochs", 100)

    # Log metrics during training
    for epoch in range(100):
        train_loss = train_one_epoch()
        mlflow.log_metric("train_loss", train_loss, step=epoch)

    # Log final metrics
    mlflow.log_metric("accuracy", 0.95)
    mlflow.log_metric("f1_score", 0.93)

    # Log model artifact
    mlflow.sklearn.log_model(model, "model")

Searching and Comparing Runs

import mlflow

# Search runs with filter
runs = mlflow.search_runs(
    experiment_names=["my-classification-experiment"],
    filter_string="metrics.accuracy > 0.9",
    order_by=["metrics.accuracy DESC"],
    max_results=10
)

# Get best run
best_run = runs.iloc[0]
print(f"Best run ID: {best_run.run_id}")
print(f"Best accuracy: {best_run['metrics.accuracy']}")

Model Registry Operations

import mlflow

# Register model from run
model_uri = f"runs:/{run_id}/model"
mlflow.register_model(model_uri, "production-classifier")

# Transition model stage
client = mlflow.tracking.MlflowClient()
client.transition_model_version_stage(
    name="production-classifier",
    version=1,
    stage="Production"
)

# Load production model
model = mlflow.pyfunc.load_model("models:/production-classifier/Production")

Integration with Babysitter SDK

Task Definition Example

const mlflowTrackingTask = defineTask({
  name: 'mlflow-experiment-tracking',
  description: 'Track ML experiment with MLflow',

  inputs: {
    experimentName: { type: 'string', required: true },
    runName: { type: 'string', required: true },
    parameters: { type: 'object', required: true },
    metrics: { type: 'object', required: true },
    modelPath: { type: 'string' }
  },

  outputs: {
    runId: { type: 'string' },
    experimentId: { type: 'string' },
    artifactUri: { type: 'string' }
  },

  async run(inputs, taskCtx) {
    return {
      kind: 'skill',
      title: `Track experiment: ${inputs.experimentName}/${inputs.runName}`,
      skill: {
        name: 'mlflow-experiment-tracker',
        context: {
          operation: 'log_run',
          experimentName: inputs.experimentName,
          runName: inputs.runName,
          parameters: inputs.parameters,
          metrics: inputs.metrics,
          modelPath: inputs.modelPath
        }
      },
      io: {
        inputJsonPath: `tasks/${taskCtx.effectId}/input.json`,
        outputJsonPath: `tasks/${taskCtx.effectId}/result.json`
      }
    };
  }
});

MCP Server Integration

Using mlflow-mcp Server

{
  "mcpServers": {
    "mlflow": {
      "command": "uvx",
      "args": ["mlflow-mcp"],
      "env": {
        "MLFLOW_TRACKING_URI": "http://localhost:5000"
      }
    }
  }
}

Available MCP Tools

• mlflow_list_experiments - List all experiments
• mlflow_search_runs - Search runs with filters
• mlflow_get_run - Get run details
• mlflow_log_metric - Log a metric
• mlflow_log_param - Log a parameter
• mlflow_list_artifacts - List run artifacts
• mlflow_get_model_version - Get model version details

Best Practices

1. Consistent Naming: Use descriptive experiment and run names
2. Complete Logging: Log all hyperparameters, not just tuned ones
3. Metric Granularity: Log metrics at appropriate intervals
4. Artifact Organization: Use consistent artifact paths
5. Model Documentation: Add descriptions to registered models
6. Stage Management: Use proper staging workflow (None -> Staging -> Production)

References

• MLflow Documentation
• MLflow MCP Server
• Official MLflow MCP (3.4+)
• MLflow Model Registry