Experiment Tracking with MLflow

Production MLflow setup for experiment tracking, model versioning, and artifact management. Covers local and remote tracking servers, model registry, and CI/CD integration.

Experiment tracking with MLflow is the practice of recording, organizing, and comparing the full lifecycle of machine learning experiments—code, hyperparameters, metrics, models, and artifacts—across multiple runs, teams, and projects. It matters when a data scientist runs dozens of training jobs, each with different configurations, and needs to answer: *Which setup worked best? Why? What changed between runs? How do I reproduce the top-performing model?* MLflow provides a standardized, open, and tool-agnostic way to answer these questions with minimal friction.

Setting Up Experiment Tracking

To begin tracking experiments, configure a backend store and a tracking URI. The tracking server can be a file-based store or a database-backed server. For a file-based store, use a directory path:

export MLFLOW_TRACKING_URI=file:///path/to/experiments

For a database-backed store, configure a PostgreSQL connection string:

export MLFLOW_TRACKING_URI=postgresql://user:password@host:port/dbname

The tracking URI must be set before any MLflow command or code execution. MLflow uses this URI to persist experiment metadata and run data.

In code, start a run with a context manager:

import mlflow

mlflow.set_experiment("image-classification-v2")

with mlflow.start_run(run_name="resnet50-finetune"):
    # Log parameters
    mlflow.log_param("learning_rate", 1e-3)
    mlflow.log_param("batch_size", 64)
    mlflow.log_param("architecture", "resnet50")

    # Log metrics
    mlflow.log_metric("train_loss", 0.42)
    mlflow.log_metric("val_accuracy", 0.87)

    # Log artifacts
    mlflow.log_artifact("model.pth")
    mlflow.log_artifact("confusion_matrix.png")

    # Log model
    mlflow.pytorch.log_model(model, "model")

This code creates a new run in the image-classification-v2 experiment and records all parameters, metrics, and artifacts. The mlflow.start_run() function creates a new run ID, which is automatically used by all subsequent log calls.

Running Experiments with Command-Line Scripts

For reproducible training scripts, use mlflow run to launch a project from a MLproject file:


name: image-classification
conda_env: conda.yaml
entry_points:
  train:
    parameters:
      learning_rate:
        type: float
        default: 0.001
      batch_size:
        type: int
        default: 32
    command: |
      python train.py --lr {learning_rate} --batch-size {batch_size}

Run the project with:

mlflow run . -P learning_rate=1e-3 -P batch_size=128

MLflow automatically installs dependencies from the conda.yaml file, runs the command, and logs all parameters, metrics, and artifacts. The command also supports --no-conda to skip environment setup and --backend to specify a remote execution backend like Docker or Kubernetes.

Logging Model Artifacts and Metadata

MLflow supports logging multiple artifact types: files, directories, models, and custom objects. The log_artifact function copies files into the run’s artifact storage, which can be local, S3, GCS, or Azure Blob.

mlflow.log_artifact("training_data.csv", artifact_path="data")
mlflow.log_artifact("model.pkl", artifact_path="models")
mlflow.log_artifact("plots/", artifact_path="visualizations")

Artifacts are stored under the run’s artifacts/ directory. The artifact_path parameter controls the subdirectory nesting.

For models, use one of the built-in model flavors:

mlflow.sklearn.log_model(model, "sklearn_model")
mlflow.xgboost.log_model(model, "xgboost_model")
mlflow.keras.log_model(model, "keras_model")

Each flavor saves the model in a standardized format with metadata like input schema, model version, and a model.yaml file. The mlflow.pytorch.log_model function, for example, saves a PyTorch model with state_dict, torchscript, and model_class.

Querying and Comparing Experiments

After running multiple experiments, retrieve and compare runs using mlflow.search_runs():

import mlflow

runs = mlflow.search_runs(
    experiment_ids=["1", "2"],
    filter_string="metrics.val_accuracy > 0.85",
    order_by=["metrics.val_accuracy DESC"],
    max_results=10
)

print(runs[["run_id", "params.learning_rate", "metrics.val_accuracy", "artifact_uri"]])

The filter_string supports comparisons, string matching, and nested field access. For example:

The returned DataFrame includes all logged parameters, metrics, and run metadata. The order_by list supports sorting on metrics, parameters, and tags.

To compare runs side-by-side, use mlflow.ui to launch a web server:

mlflow ui --host 0.0.0.0 --port 5000

The UI displays a table of runs with columns for parameters, metrics, tags, and artifacts. Clicking a run shows a detailed view with logs, model artifacts, and visualizations.

Configuring the Tracking Server

For shared, persistent tracking, run an MLflow tracking server:

mlflow server \
  --host 0.0.0.0 \
  --port 5000 \
  --backend-store-uri postgresql://user:pass@host:5432/mlflow \
  --artifact-root s3://my-bucket/mlflow-artifacts \
  --workers 4

The server exposes an HTTP API at /api/2.0/mlflow/. The /runs/search endpoint accepts a JSON body:

{
  "experiment_ids": ["1"],
  "filter_string": "metrics.accuracy > 0.9",
  "max_results": 50,
  "order_by": ["metrics.val_accuracy DESC"],
  "view_type": "ACTIVE_ONLY"
}

Response includes a list of runs, each with run_id, experiment_id, metrics, params, tags, and artifact_uri.

Common failure modes include:

Managing Model Lifecycle and Versioning

After training, register models using mlflow.register_model:

mlflow.register_model(
    model_uri="runs:/1234567890/outputs",
    name="customer-churn-predictor"
)

This creates a model in the model registry with a version number. The model registry supports stages: Staging, Production, Archived. Promote a model:

mlflow.register_model(
    model_uri="runs:/1234567890/outputs",
    name="customer-churn-predictor"
)

mlflow.set_model_version_tag(
    name="customer-churn-predictor",
    version=1,
    key="deployed_by",
    value="[email protected]"
)

Versioning is automatic: each mlflow.log_model call creates a new version. The model registry tracks the model lineage, including the run that produced it, the parameters, and the artifact URI.

To load a model:

model = mlflow.pyfunc.load_model("models:/customer-churn-predictor/1")
predictions = model.predict(data)

The mlflow.pyfunc.load_model function resolves the model from the registry, downloads the artifact, and reconstructs the model object.

Advanced: Custom Tags and Metadata

Tags are key-value pairs attached to runs, experiments, and models. Use mlflow.set_tag to add metadata:

mlflow.set_tag("user", "alice")
mlflow.set_tag("training_set", "2023-01-01-to-2023-06-30")
mlflow.set_tag("deployed", "true")

Tags appear in the tracking UI and are searchable. They are useful for categorizing runs by environment, dataset, or deployment context.

For complex metadata, use mlflow.set_tags with a dictionary:

mlflow.set_tags({
    "user": "alice",
    "training_set": "2023-01-01-to-2023-06-30",
    "deployed": True,
    "description": "Fine-tuned ResNet50 on medical images"
})

Tags are stored in the backend store and can be used in filter_string queries.

Common Pitfalls and Silent Failures

Automating Experiment Workflows

Automate experiment tracking using shell scripts or CI/CD pipelines. In a Makefile:

experiment:
	mlflow run . \
		-P learning_rate=1e-3 \
		-P batch_size=64 \
		--experiment-name "image-classification" \
		--no-conda

In a GitHub Actions workflow:

- name: Run MLflow Experiment
  run: |
    export MLFLOW_TRACKING_URI=postgresql://user:pass@host:5432/mlflow
    mlflow run . -P learning_rate=1e-4 -P epochs=100

The tracking URI is shared across jobs, so runs from different workflows appear in the same experiment. Use mlflow.set_experiment to group related runs under a single name.

Summary

MLflow experiment tracking is the foundation for reproducible and scalable machine learning. It unifies code, hyperparameters, metrics, models, and artifacts into a single, queryable record. The key is consistent use of mlflow.start_run, mlflow.log_param, mlflow.log_metric, and mlflow.log_artifact. The tracking server, model registry, and command-line interface form a complete ecosystem for managing experiments across teams and projects. When experiments are logged well, the next step—comparing, deploying, and maintaining models—becomes predictable, repeatable, and collaborative.

This page was rewritten on 17 October 2026. It replaced a templated version that was not specific to its own title. How these pages are written and checked. If anything here is wrong, tell us at [email protected] and we will correct it.