MLflow Tracking

The MLflow Tracking is an API and UI for logging parameters, code versions, metrics, and output files when running your machine learning code and for later visualizing the results. MLflow Tracking provides Python, REST, R, and Java APIs.

_images/tracking-metrics-ui-temp.png

A screenshot of the MLflow Tracking UI, showing a plot of validation loss metrics during model training.

Quickstart

If you haven’t used MLflow Tracking before, we strongly recommend going through the following quickstart tutorial.

Concepts

Runs

MLflow Tracking is organized around the concept of runs, which are executions of some piece of data science code, for example, a single python train.py execution. Each run records metadata (various information about your run such as metrics, parameters, start and end times) and artifacts (output files from the run such as model weights, images, etc).

Experiments

An experiment groups together runs for a specific task. You can create an experiment using the CLI, API, or UI. The MLflow API and UI also let you create and search for experiments. See Organizing Runs into Experiments for more details on how to organize your runs into experiments.

Tracking Runs

MLflow Tracking APIs provide a set of functions to track your runs. For example, you can call mlflow.start_run() to start a new run, then call Logging Functions such as mlflow.log_param() and mlflow.log_metric() to log a parameters and metrics respectively. Please visit the Tracking API documentation for more details about using these APIs.

import mlflow

with mlflow.start_run():
    mlflow.log_param("lr", 0.001)
    # Your ml code
    ...
    mlflow.log_metric("val_loss", val_loss)

Alternatively, Auto-logging offers the ultra-quick setup for starting MLflow tracking. This powerful feature allows you to log metrics, parameters, and models without the need for explicit log statements - all you need to do is call mlflow.autolog() before your training code. Auto-logging supports popular libraries such as Scikit-learn, XGBoost, PyTorch, Keras, Spark, and more. See Automatic Logging Documentation for supported libraries and how to use auto-logging APIs with each of them.

import mlflow

mlflow.autolog()

# Your training code...

Note

By default, without any particular server/database configuration, MLflow Tracking logs data to the local mlruns directory. If you want to log your runs to a different location, such as a remote database and cloud storage, to share your results with your team, follow the instructions in the Set up MLflow Tracking Environment section.

Tracking Datasets

MLflow offers the ability to track datasets that are associated with model training events. These metadata associated with the Dataset can be stored through the use of the mlflow.log_input() API. To learn more, please visit the MLflow data documentation to see the features available in this API.

Explore Runs and Results

Tracking UI

The Tracking UI lets you visually explore your experiments and runs, as shown on top of this page.

  • Experiment-based run listing and comparison (including run comparison across multiple experiments)

  • Searching for runs by parameter or metric value

  • Visualizing run metrics

  • Downloading run results (artifacts and metadata)

If you log runs to a local mlruns directory, run the following command in the directory above it, then access http://127.0.0.1:5000 in your browser.

mlflow ui --port 5000

Alternatively, the MLflow Tracking Server serves the same UI and enables remote storage of run artifacts. In that case, you can view the UI at http://<IP address of your MLflow tracking server>:5000 from any machine that can connect to your tracking server.

Querying Runs Programmatically

You can also access all of the functions in the Tracking UI programmatically with MlflowClient.

For example, the following code snippet search for runs that has the best validation loss among all runs in the experiment.

client = mlflow.tracking.MlflowClient()
experiment_id = "0"
best_run = client.search_runs(
    experiment_id, order_by=["metrics.val_loss ASC"], max_results=1
)[0]
print(best_run.info)
# {'run_id': '...', 'metrics': {'val_loss': 0.123}, ...}

Set up the MLflow Tracking Environment

Note

If you just want to log your experiment data and models to local files, you can skip this section.

MLflow Tracking supports many different scenarios for your development workflow. This section will guide you through how to set up the MLflow Tracking environment for your particular use case. From a bird’s-eye view, the MLflow Tracking environment consists of the following components.

Components

MLflow Tracking APIs

You can call MLflow Tracking APIs in your ML code to log runs and communicate with the MLflow Tracking Server if necessary.

Backend Store

The backend store persists various metadata for each Run, such as run ID, start and end times, parameters, metrics, etc. MLflow supports two types of storage for the backend: file-system-based like local files and database-based like PostgreSQL.

Artifact Store

Artifact store persists (typicaly large) arifacts for each run, such as model weights (e.g. a pickled scikit-learn model), images (e.g. PNGs), model and data files (e.g. Parquet file). MLflow stores artifacts ina a local file (mlruns) by default, but also supports different storage options such as Amazon S3 and Azure Blob Storage.

MLflow Tracking Server (Optional)

MLflow Tracking Server is a stand-alone HTTP server that provides REST APIs for accessing backend and/or artifact store. Tracking server also offers flexibility to configure what data to server, govern access control, versioning, and etc. Read MLflow Tracking Server documentation for more details.

Common Setups

By configuring these components properly, you can create an MLflow Tracking environment suitable for your team’s development workflow. The following diagram and table show a few common setups for the MLflow Tracking environment.

_images/tracking-setup-overview.png

1. Localhost (default)

2. Local Tracking with Local Database

3. Remote Tracking with MLflow Tracking Server

Scenario

Solo development

Solo development

Team development

Use Case

By default, MLflow records metadata and artifacts for each run to a local directory, mlruns. This is the simplest way to get started with MLflow Tracking, without setting up any external server, database, and storage.

The MLflow client can interface with a SQLAlchemy-compatible database (e.g., SQLite, PostgreSQL, MySQL) for the backend. Saving metadata to a database allows you cleaner management of your experiment data while skipping the effort of setting up a server.

MLflow Tracking Server can be configured with an artifacts HTTP proxy, passing artifact requests through the tracking server to store and retrieve artifacts without having to interact with underlying object store services. This is particularly useful for team development scenarios where you want to store artifacts and experiment metadata in a shared location with proper access control.

Tutorial

QuickStart

Tracking Experiments using a Local Database

Remote Experiment Tracking with MLflow Tracking Server

Other Configuration with MLflow Tracking Server

MLflow Tracking Server provides customizability for other special use cases. Please follow Remote Experiment Tracking with MLflow Tracking Server for learning the basic setup and continue to the following materials for advanced configurations to meet your needs.

Using MLflow Tracking Server Locally

You can of course run MLflow Tracking Server locally. While this doesn't provide much additional benefit over directly using the local files or database, might useful for testing your team development workflow locally or running your machine learning code on a container environment.

FAQ

Can I launch multiple runs in parallel?

Yes, MLflow supports launching multiple runs in parallel e.g. multi processing / threading. See Launching Multiple Runs in One Program for more details.

How can I organize many MLflow Runs neatly?

MLflow provides a few ways to organize your runs:

  • Organize runs into experiments - Experiments are logical containers for your runs. You can create an experiment using the CLI, API, or UI.

  • Create child runs - You can create child runs under a single parent run to group them together. For example, you can create a child run for each fold in a cross-validation experiment.

  • Add tags to runs - You can associate arbitrary tags with each run, which allows you to filter and search runs based on tags.

Can I directly access remote storage without running the Tracking Server?

Yes, while it is best practice to have the MLflow Tracking Server as a proxy for artifacts access for team development workflows, you may not need that if you are using it for personal projects or testing. You can achieve this by following the workaround below:

  1. Set up artifacts configuration such as credentials and endpoints, just like you would for the MLflow Tracking Server. See configure artifact storage for more details.

  2. Create an experiment with an explicit artifact location,

experiment_name = "your_experiment_name"
mlflow.create_experiment(experiment_name, artifact_location="s3://your-bucket")
mlflow.set_experiment(experiment_name)

Your runs under this experiment will log artifacts to the remote storage directly.

How to integrate MLflow Tracking with Model Registry?

To use the Model Registry functionality with MLflow tracking, you must use database backed store such as PostgresQL and log a model using the log_model methods of the corresponding model flavors. Once a model has been logged, you can add, modify, update, or delete the model in the Model Registry through the UI or the API. See Backend Stores and Common Setups for how to configures backend store properly for your workflow.

How to include additional decription texts about the run?

A system tag mlflow.note.content can be used to add descriptive note about this run. While the other system tags are set automatically, this tag is not set by default and users can override it to include additional information about the run. The content will be displayed on the run’s page under the Notes section.